import os
import csv
import gradio as gr
import tensorflow as tf
import numpy as np
import pandas as pd
from datetime import datetime
import utils
from huggingface_hub import Repository
import itertools
import GPyOpt

# Unique phase elements

# Load access tokens
WRITE_TOKEN = os.environ.get("WRITE_PER") # write

# Logs repo path
dataset_url = "https://huggingface.co/datasets/sandl/upload_alloy_hardness"
dataset_path = "logs_alloy_hardness.csv"

scaling_factors = {'PROPERTY: Calculated Density (g/cm$^3$)': (5.5, 13.7),
                   'PROPERTY: Calculated Young modulus (GPa)': (77.0, 336.0),
                   'PROPERTY: HV': (107.0, 1183.0),
                   'PROPERTY: YS (MPa)': (62.0, 3416.0)}

input_mapping = {'PROPERTY: BCC/FCC/other': {'BCC': 0, 'FCC': 1, 'OTHER': 2},#, 'nan': 2},
                 'PROPERTY: Processing method': {'ANNEAL': 0, 'CAST': 1, 'OTHER': 2, 'POWDER': 3, 'WROUGHT': 4},#, 'nan': 2},
                 'PROPERTY: Microstructure': {'B2': 0, 'B2+BCC': 1, 'B2+L12': 2, 'B2+Laves+Sec.': 3, 'B2+Sec.': 4, 'BCC': 5,
                                              'BCC+B2': 6, 'BCC+B2+FCC': 7, 'BCC+B2+FCC+Sec.': 8, 'BCC+B2+L12': 9, 'BCC+B2+Laves': 10,
                                              'BCC+B2+Sec.': 11, 'BCC+BCC': 12, 'BCC+BCC+HCP': 13, 'BCC+BCC+Laves': 14,
                                              'BCC+BCC+Laves(C14)': 15, 'BCC+BCC+Laves(C15)': 16, 'BCC+FCC': 17, 'BCC+HCP': 18,
                                              'BCC+Laves': 19, 'BCC+Laves(C14)': 20, 'BCC+Laves(C15)': 21, 'BCC+Laves+Sec.': 22,
                                              'BCC+Sec.': 23, 'FCC': 24, 'FCC+B2': 25, 'FCC+B2+Sec.': 26, 'FCC+BCC': 27,
                                              'FCC+BCC+B2': 28, 'FCC+BCC+B2+Sec.': 29, 'FCC+BCC+BCC': 30, 'FCC+BCC+Sec.': 31,
                                              'FCC+FCC': 32, 'FCC+HCP': 33, 'FCC+HCP+Sec.': 34, 'FCC+L12': 35, 'FCC+L12+B2': 36,
                                              'FCC+L12+Sec.': 37, 'FCC+Laves': 38, 'FCC+Laves(C14)': 39, 'FCC+Laves+Sec.': 40,
                                              'FCC+Sec.': 41, 'L12+B2': 42, 'Laves(C14)+Sec.': 43, 'OTHER': 44},#, 'nan': 44},
                 'PROPERTY: Single/Multiphase': {'': 0, 'M': 1, 'S': 2, 'OTHER': 3}}#, 'nan': 3}}

unique_phase_elements = ['B2', 'BCC', 'FCC', 'HCP', 'L12', 'Laves', 'Laves(C14)', 'Laves(C15)', 'Sec.', 'OTHER']

input_cols = {
    "PROPERTY: Alloy formula": "(PROPERTY: Alloy formula) "
                               "Enter alloy formula using proportions representation (i.e. Al0.25 Co1 Fe1 Ni1)",
    "PROPERTY: Single/Multiphase": "(PROPERTY: Single/Multiphase) "
                                   "Choose between Single (S), Multiphase (M) and other (OTHER)",
    "PROPERTY: BCC/FCC/other": "(PROPERTY: BCC/FCC/other) "
                               "Choose between BCC, FCC and other ",
    "PROPERTY: Processing method": "(PROPERTY: Processing method) "
                                   "Choose your processing method (ANNEAL, CAST, POWDER, WROUGHT or OTHER)",
    "PROPERTY: Microstructure": "(PROPERTY: Microstructure) "
                                "Choose the microstructure (SEC means the secondary/tertiary microstructure is not one of FCC, BCC, HCP, L12, B2, Laves, Laves (C14), Laves (C15))",
}

def process_microstructure(list_phases):
    permutations = list(itertools.permutations(list_phases))
    permutations_strings = [str('+'.join(list(e))) for e in permutations]
    for e in permutations_strings:
        if e in list(input_mapping['PROPERTY: Microstructure'].keys()):
            return e
    return 'OTHER'
    
def write_logs(message, message_type="Prediction"):
    """
    Write logs
    """
    #with Repository(local_dir="data", clone_from=dataset_url, use_auth_token=WRITE_TOKEN).commit(commit_message="from private", blocking=False):
     #   with open(dataset_path, "a") as csvfile:
      #          writer = csv.DictWriter(csvfile, fieldnames=["name", "message", "time"])
       #         writer.writerow(
        #            {"name": message_type, "message": message, "time": str(datetime.now())}
         #       )
    return 
    
def predict(x, request: gr.Request):
    """
    Predict the hardness and yield strength using the ML model. Input data is a dataframe
    """
    loaded_model = tf.keras.models.load_model("hardness_nn_graph_separate_elements.h5")
    print("summary is", loaded_model.summary())
    #x = x.replace("", 0)
    x = np.asarray(x).astype("float32")
    y = loaded_model.predict(x)
    y_hardness = y[0][0]
    y_ys = y[0][1]
    minimum_hardness, maximum_hardness = scaling_factors['PROPERTY: HV']
    minimum_ys, maximum_ys = scaling_factors['PROPERTY: YS (MPa)']
    print("Prediction is ", y)
    if request is not None:   # Verify if request is not None (when building the app the first request is None)
        message = f"{request.username}_{request.client.host}"
        print("MESSAGE")
        print(message)
        res = write_logs(message)
    #interpret_fig = utils.interpret(x)
    return (round(y_hardness*(maximum_hardness-minimum_hardness)+minimum_hardness, 2), 12,
            round(y_ys*(maximum_ys-minimum_ys)+minimum_ys, 2), 12)

def fit_outputs_constraints(x, hardness_target, ys_target, request: gr.Request):
    print("*********")
    print(len(x))
    print(x)
    print(len(x[0]))
    predictions = predict(x, request)
    error_hardness = np.sqrt(np.square(predictions[0]-float(hardness_target)))
    error_ys = np.sqrt(np.square(predictions[2]-float(ys_target)))
    print("Optimization step is ", predictions, float(hardness_target), float(ys_target),
          error_hardness, error_ys)
    return error_hardness + error_ys

def predict_inverse(hardness_original_target, ys_original_target, metals_to_use, request: gr.Request):

    one_hot_columns = utils.return_feature_names()
    min_df_hardness, max_df_hardness = scaling_factors["PROPERTY: HV"]
    hardness_original_target = float(hardness_original_target)
    min_df_ys, max_df_ys = scaling_factors["PROPERTY: YS (MPa)"]
    ys_original_target = float(ys_original_target)

    hardness_target = (hardness_original_target-min_df_hardness)/(max_df_hardness-min_df_hardness)
    ys_target = (ys_original_target-min_df_ys)/(max_df_ys-min_df_ys)

    continuous_variables = ['PROPERTY: Calculated Density (g/cm$^3$)',
                        'PROPERTY: Calculated Young modulus (GPa)',
                        'PROPERTY: Metal Al', 'PROPERTY: Metal Co',
                        'PROPERTY: Metal Fe', 'PROPERTY: Metal Ni', 'PROPERTY: Metal Si',
                        'PROPERTY: Metal Cr', 'PROPERTY: Metal Nb', 'PROPERTY: Metal Ti',
                        'PROPERTY: Metal Mn', 'PROPERTY: Metal V', 'PROPERTY: Metal Mo',
                        'PROPERTY: Metal Cu', 'PROPERTY: Metal Ta', 'PROPERTY: Metal Zr',
                        'PROPERTY: Metal Hf', 'PROPERTY: Metal W', 'PROPERTY: Metal Zn',
                        'PROPERTY: Metal Sn', 'PROPERTY: Metal Re', 'PROPERTY: Metal C',
                        'PROPERTY: Metal Pd', 'PROPERTY: Metal Sc', 'PROPERTY: Metal Y']
    categorical_variables = list(one_hot_columns)
    for c in continuous_variables:
        categorical_variables.remove(c)

    # Metals constraints 
    metals_elements = [c for c in continuous_variables if c.startswith("PROPERTY: Metal")]
    # metals_to_use = ['Al', 'Co', 'Fe', 'Cr']
    metals_to_use = ["PROPERTY: Metal " + metals_to_use[i] for i in range(len(metals_to_use))]

    # Domain 
    domain = []
    for c in one_hot_columns:
        if c in continuous_variables:
            if c.startswith("PROPERTY: Metal") and c not in metals_to_use:
                domain.append({'name': str(c), 'type': 'continuous', 'domain': (0.0, 0.)})
            else:
                domain.append({'name': str(c), 'type': 'continuous', 'domain': (0.1, 1.)})#(0.,1.)})
        else:
            domain.append({'name': str(c), 'type': 'discrete', 'domain': (0,1)})

    print("************")
    print("Domain")
    print(domain)
    # Constraints 
    constraints = []
    constrained_columns = ['Single/Multiphase', 'Preprocessing method', 'BCC/FCC/other'] #'PROPERTY: Metal']#, 'Microstructure']
    for constraint in constrained_columns:
        sum_string = ''
        for i in range (len(one_hot_columns)):
            column_one_hot = one_hot_columns[i]
            if column_one_hot.startswith(constraint):
                sum_string = sum_string+"+x[:," + str(i) + "]"
        constraints.append({'name': constraint + "+1", 'constraint': sum_string + '-1'})
        constraints.append({'name': constraint + "-1", 'constraint': '-1*(' + sum_string + ')+1'})
    print("********************")
    print("Constraints")
    print(constraints)

    def fit_outputs(x):
        return fit_outputs_constraints(x, hardness_target, ys_target, request)

    opt = GPyOpt.methods.BayesianOptimization(f = fit_outputs,            # function to optimize       
                                              domain = domain,        # box-constraints of the problem
                                              constraints = constraints,
                                              acquisition_type ='LCB',       # LCB acquisition
                                              acquisition_weight = 0.1)   # Exploration exploitation
    # it may take a few seconds
    opt.run_optimization(max_iter=5)
    # opt.plot_convergence()
    x_best = opt.X[np.argmin(opt.Y)]
    best_params = dict(zip(
        [el['name'] for el in domain],
        [[x] for x in x_best]))
    optimized_x = pd.DataFrame.from_dict(best_params)
    print("Optimized parameters")
    print(optimized_x.columns)
    for c in optimized_x.columns:
        if c in continuous_variables:
            if c in ['PROPERTY: Calculated Density (g/cm$^3$)', 'PROPERTY: Calculated Young modulus (GPa)']:
                print(c)
                print(optimized_x[c])
                optimized_x[c]=round(optimized_x[c]*(scaling_factors[c][1]-scaling_factors[c][0])+scaling_factors[c][0], 2)
    result = optimized_x
    result = result[result>0.0].dropna(axis=1)

    # Normalize metals outputs
    sum_metals = np.sum(result[c] for c in list(result.columns) if c.startswith("PROPERTY: Metal"))
    for column in result.columns:
        if column.startswith("PROPERTY: Metal"):
            result[column]/= sum_metals
            result[column] = round(result[column], 2)

    columns = list(result.columns)
    print("cccccccccc")
    print(columns)

    print((result[columns[2:-3]], columns[-3], result.at[0, columns[0]], result.at[0, columns[1]], columns[-2], columns[-1]))
    print("_")
    print(result[columns[2:-3]])
    print("_")
    print(columns[-3])
    print("_")
    print(result.at[0, columns[0]])
    print("_")
    print(result.at[0, columns[1]])
    print("_")
    print(columns[-2], columns[-1])
    return (result[columns[2:-3]], columns[-3], result.at[0, columns[0]],
            result.at[0, columns[1]], columns[-2], columns[-1])


example_inputs = [820, 1800, ['Al', 'Fe', 'Ni']]

css_styling = """#submit {background: #1eccd8} 
#submit:hover {background: #a2f1f6} 
.output-image, .input-image, .image-preview {height: 250px !important}
.output-plot {height: 250px !important}"""

light_theme_colors = gr.themes.Color(c50="#e4f3fa", # Dataframe background cell content - light mode only
                                c100="#e4f3fa", # Top corner of clear button in light mode + markdown text in dark mode
                                c200="#a1c6db", # Component borders
                                c300="#FFFFFF", # 
                                c400="#e4f3fa", # Footer text
                                c500="#0c1538", # Text of component headers in light mode only
                                c600="#a1c6db", # Top corner of button in dark mode
                                c700="#475383", # Button text in light mode + component borders in dark mode
                                c800="#0c1538", # Markdown text in light mode
                                c900="#a1c6db", # Background of dataframe - dark mode
                                c950="#0c1538") # Background in dark mode only
# secondary color used for highlight box content when typing in light mode, and download option in dark mode
# primary color used for login button in dark mode
osium_theme = gr.themes.Default(primary_hue="cyan", secondary_hue="cyan", neutral_hue=light_theme_colors)
page_title = "Alloys' hardness and yield strength prediction"
favicon_path = "osiumai_favicon.ico"
logo_path  = "osiumai_logo.jpg"
html = f"""<html> <link rel="icon" type="image/x-icon" href="file={favicon_path}">
<img src='file={logo_path}' alt='Osium AI logo' width='200' height='100'> </html>"""


with gr.Blocks(css=css_styling, title=page_title, theme=osium_theme) as demo:
    #gr.HTML(html)
    gr.Markdown("# <p style='text-align: center;'>Get optimal alloy recommendations based on your target performance</p>")
    gr.Markdown("This AI model provides a recommended alloy formula, microstructure and processing conditions based on your target hardness and yield strength")
    with gr.Row():
        clear_button = gr.Button("Clear")
        prediction_button = gr.Button("Predict", elem_id="submit")
    with gr.Row():
        with gr.Column():
            gr.Markdown("### The target performance of your alloy")
            input_hardness = gr.Text(label="Enter your target hardness (in HV)")
            input_yield_strength = gr.Text(label="Enter your target yield strength (MPa)")
            gr.Markdown('### Your metallic elements constraints')
            metals_constraints = gr.CheckboxGroup(
                choices=['Al', 'Co', 'Fe', 'Ni', 'Si', 'Cr', 'Nb', 'Ti',
                        'Mn', 'V', 'Mo', 'Cu', 'Ta', 'Zr',
                        'Hf', 'W', 'Zn', 'Sn', 'Re', 'C',
                        'Pd', 'Sc', 'Y'], label="Your metals constraints",
            )
        with gr.Column():
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Your optimal alloy formula and processing conditions")
                    optimal_formula = gr.DataFrame(label="Your optimal alloy formula", wrap=True)
                    optimal_processing_method = gr.Text(label="Processing method")
                    gr.Markdown("### Additional information about your optimal alloy")
                    density = gr.Text(label="Density (g/cm3)")
                    young_modulus = gr.Text(label = "Young modulus (GPa)")
                    microstructure = gr.Text(label="Microstructure (BCC/FCC/Other)")
                    phase = gr.Text(label="Number of phases (S/M)")

    with gr.Row():
        gr.Examples([example_inputs], [input_hardness, input_yield_strength, metals_constraints])
            
            

    prediction_button.click(
        fn=predict_inverse,
        inputs=[input_hardness, input_yield_strength, metals_constraints],
        outputs=[optimal_formula, optimal_processing_method, density, young_modulus, microstructure, phase],
        show_progress=True,
    )
    clear_button.click(
        lambda x: [gr.update(value=None)] * 9,
        [],
        [
            input_hardness,
            input_yield_strength,
            metals_constraints, 
            optimal_formula,
            optimal_processing_method,
            density, young_modulus,
            microstructure, phase
        ],
    )


if __name__ == "__main__":
    demo.queue(concurrency_count=2)
    demo.launch()