import keras
from keras.layers import Dense, Dropout, Input,BatchNormalization
import tensorflow as tf
import h5py
import numpy as np
import pandas as pd
import time
import sklearn.model_selection as sk
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.pipeline import Pipeline, make_pipeline
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from imblearn.over_sampling import SMOTE
from imblearn.under_sampling import RandomUnderSampler
from imblearn.over_sampling import RandomOverSampler
from collections import Counter
import os
from preprocessing import *
from sklearn import set_config
from importlib import reload
set_config(transform_output = "pandas")
import optuna
import pickle

def objective(trial, X, y, species_columns):
    
    model_type = trial.suggest_categorical("model", ["small", "large", "paper"])
    scaler_type = trial.suggest_categorical("scaler", ["standard", "minmax"])
    sampling_type = trial.suggest_categorical("sampling", ["over", "off"])
    loss_variant = trial.suggest_categorical("loss", ["huber", "huber_mass_balance"])
    delta = trial.suggest_float("delta", 0.5, 5.0)
    
    preprocess = preprocessing()
    X, y = preprocess.cluster(df_design[species_columns], df_results[species_columns])
    X_train, X_test, y_train, y_test = preprocess.split(X, y, ratio = 0.2)
    X_train, y_train = preprocess.balancer(X_train, y_train, strategy = sampling_type)
    preprocess.scale_fit(X_train, y_train, scaling = "global", type=scaler_type)
    X_train, X_test, y_train, y_test = preprocess.scale_transform(X_train, X_test, y_train, y_test)
    X_train, X_val, y_train, y_val = preprocess.split(X_train, y_train, ratio = 0.1)

    column_dict = {"Ba": X.columns.get_loc("Ba"), "Barite": X.columns.get_loc("Barite"), "Sr": X.columns.get_loc(
        "Sr"), "Celestite": X.columns.get_loc("Celestite"), "H": X.columns.get_loc("H"), "H": X.columns.get_loc("H"), "O": X.columns.get_loc("O")}

    h1 = trial.suggest_float("h1", 0.1, 1.0)
    h2 = trial.suggest_float("h2", 0.1, 1.0)
    h3 = trial.suggest_float("h3", 0.1, 1.0)

    model = model_definition(model_type)
    
    lr_schedule = keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate=0.001,
        decay_steps=2000,
        decay_rate=0.9,
        staircase=True
    )
    optimizer = keras.optimizers.Adam(learning_rate=lr_schedule)
    
    model.compile(optimizer=optimizer, loss=custom_loss(preprocess, column_dict, h1, h2, h3, scaler_type, loss_variant, delta),
                  metrics=[huber_metric(preprocess, scaler_type, delta), mass_balance_metric(preprocess, column_dict, scaler_type)])
    
    callback = keras.callbacks.EarlyStopping(monitor='loss', patience=3)
    history = model.fit(X_train.loc[:, X_train.columns != "Class"], 
                        y_train.loc[:, y_train.columns != "Class"], 
                        batch_size=512, 
                        epochs=100, 
                        validation_data=(X_val.loc[:, X_val.columns != "Class"], y_val.loc[:, y_val.columns != "Class"]),
                        callbacks=[callback])
    
    prediction_huber_overall = model.evaluate(
        X_test.loc[:, X_test.columns != "Class"], y_test.loc[:, y_test.columns != "Class"])[1]
    prediction_huber_non_reactive = model.evaluate(
        X_test[X_test['Class'] == 0].iloc[:, X_test.columns != "Class"], y_test[X_test['Class'] == 0].iloc[:, y_test.columns != "Class"])[1]
    prediction_huber_reactive = model.evaluate(
        X_test[X_test['Class'] == 1].iloc[:, X_test.columns != "Class"], y_test[X_test['Class'] == 1].iloc[:, y_test.columns != "Class"])[1]
    prediction_mass_balance_overall = model.evaluate(
        X_test.loc[:, X_test.columns != "Class"], y_test.loc[:, y_test.columns != "Class"])[2]
    prediction_mass_balance_non_reactive = model.evaluate(
        X_test[X_test['Class'] == 0].iloc[:, X_test.columns != "Class"], y_test[X_test['Class'] == 0].iloc[:, y_test.columns != "Class"])[2]
    prediction_mass_balance_reactive = model.evaluate(
        X_test[X_test['Class'] == 1].iloc[:, X_test.columns != "Class"], y_test[X_test['Class'] == 1].iloc[:, y_test.columns != "Class"])[2]
    
    mass_balance_results = mass_balance_evaluation(model, X_test, preprocess)
    mass_balance_ratio = len(mass_balance_results[mass_balance_results < 1e-5]) / len(mass_balance_results)
    
    results_save_path = os.path.join("./results/", "results.csv")
    results_df = pd.DataFrame({
        "trial": [trial.number],
        "prediction_huber_overall": [prediction_huber_overall],
        "prediction_huber_non_reactive": [prediction_huber_non_reactive],
        "prediction_huber_reactive": [prediction_huber_reactive],
        "prediction_mass_balance_overall": [prediction_mass_balance_overall],
        "prediction_mass_balance_non_reactive": [prediction_mass_balance_non_reactive],
        "prediction_mass_balance_reactive": [prediction_mass_balance_reactive]
    })
    
    if not os.path.isfile(results_save_path):
        results_df.to_csv(results_save_path, index=False)
    else:
        results_df.to_csv(results_save_path, mode='a', header=False, index=False)
    
    
    model_save_path_trial = os.path.join("./results/models/", f"model_trial_{trial.number}.keras")
    history_save_path_trial = os.path.join("./results/history/", f"history_trial_{trial.number}.pkl")

    model.save(model_save_path_trial)
    with open(history_save_path_trial, 'wb') as f:
        pickle.dump(history.history, f)

    return prediction_huber_overall, mass_balance_ratio

if __name__ == "__main__":

    
    print(os.path.abspath("./datasets/barite_50_4_corner.h5"))
    data_file = h5py.File("./datasets/barite_50_4_corner.h5")
    design = data_file["design"]
    results = data_file["result"]

    df_design = pd.DataFrame(np.array(design["data"]).transpose(), columns = np.array(design["names"].asstr()))
    df_results = pd.DataFrame(np.array(results["data"]).transpose(), columns = np.array(results["names"].asstr()))

    data_file.close()

    species_columns = ['H', 'O', 'Ba', 'Cl', 'S', 'Sr', 'Barite', 'Celestite']
    
    study = optuna.create_study(storage="sqlite:///model_large_optimization.db", study_name="model_optimization", directions=["minimize", "maximize"])
    study.optimize(lambda trial: objective(trial, df_design, df_results, species_columns), n_trials=1000)

    print("Number of finished trials: ", len(study.trials))

    print("Best trial:")
    trial = study.best_trial

    print("  Value: ", trial.value)

    print("  Params: ")
    for key, value in trial.params.items():
        print("    {}: {}".format(key, value))