
import xgboost as xgb
from sklearn.metrics import matthews_corrcoef, accuracy_score
import optuna
from utils import Dataset
import pandas  as pd



def objective(trial,dataset:Dataset,num_boost_round:int)->float:
    """function to maximize during the tuning phase

    Args:
        trial (??): optuna stuff
        dataset (Dataset): dataset to use (containing train  and validation)
        num_boost_round (int): number of iteration of xgboost

    Returns:
        float: validation MCC
    """
    #These are the parameters usually used
    params = dict(
                learning_rate = trial.suggest_float("learning_rate", 0.01, 0.2),
                max_depth= trial.suggest_int("max_depth",5, 15),
                min_child_weight = trial.suggest_int("min_child_weight", 1, 8),
                gamma = trial.suggest_float("gamma", 0, 10),
                subsample = trial.suggest_float("subsample", 0.01,1),
                colsample_bytree = trial.suggest_float("colsample_bytree", 0.01,1),
                alpha = trial.suggest_float("alpha", 0, 10),
                objective= 'multi:softprob', 
                nthread=4, 
                 num_class= 5,
                seed=27)
    params['lambda'] = trial.suggest_float("lambda", 0, 10)

  
    dtrain = xgb.DMatrix(dataset.X_train,dataset.y_train, 
                         enable_categorical=True,
                         weight=dataset.weight_train)
    dvalid = xgb.DMatrix(dataset.X_valid,dataset.y_valid,
                         enable_categorical=True,
                         )


    bst = xgb.train(params, dtrain,verbose_eval=False, num_boost_round=num_boost_round,
                    evals = [(dtrain, "train"), (dvalid, "valid")],
                    early_stopping_rounds=100)

    preds = bst.predict(dvalid)
    ##MCC is more solid
    mcc = matthews_corrcoef(dataset.y_valid,preds.argmax(1))    
    
    return mcc


def train(dataset:Dataset,n_trials:int=1000,timeout:int=600,num_boost_round:int=600)->(xgb.Booster, dict):
    """optuna search procedure

    Args:
        dataset (Dataset): dataset to use (containing train  and validation)
        n_trials (int, optional): number of combination to try. Defaults to 1000.
        timeout (int, optional): maximum time before stopping. Defaults to 600.
        num_boost_round (int, optional): number of iteration of a single boost model. Defaults to 600.

    Returns:
        trained xgboost and a dictionary containing the best parameters
    """


    study = optuna.create_study(direction="maximize")
    study.optimize(lambda trial: objective(trial,dataset,num_boost_round), n_trials=n_trials, timeout=timeout)

    params_final = dict(
                objective= 'multi:softprob', 
                nthread=4, 
                num_class= 5, ## 5 classi india
                seed=27)
    params_final.update(study.best_params)
    dtrain = xgb.DMatrix(dataset.X_train,dataset.y_train, 
                        enable_categorical=True,
                        weight=dataset.weight_train)
    dvalid = xgb.DMatrix(dataset.X_valid,dataset.y_valid,
                        enable_categorical=True,
                        )   
    bst = xgb.train(params_final, dtrain,verbose_eval=False, num_boost_round=num_boost_round,
                    evals = [(dtrain, "train"), (dvalid, "valid")],
                    early_stopping_rounds=100,)
    return bst,params_final, study


def gain_accuracy_train(dataset:Dataset,feat_imp:pd.DataFrame,num_boost_round:int=600,params:dict={})->(pd.DataFrame,xgb.Booster,int):
    """Starting from the most important feature, add one feature, train the model and get mcc and acc on the validation

    Args:
        dataset (Dataset):  dataset to use (containing train  and validation)
        feat_imp (pd.DataFrame): feature importance dataset computed using feat_imp = pd.Series(best_model.get_fscore()).sort_values(ascending=False)
        num_boost_round (int): number of iteration of a single boost model. Defaults to 600.
        params (dict): dictionary of best parameters returned from the function `train`

    Returns:
       dataframe with N-variables, ACC, MCC for each N 
    """

    tot = []
    for i in range(1,dataset.X_train.shape[1]):
        
        dtrain_FS = xgb.DMatrix(dataset.X_train[list(feat_imp.head(i).index)],dataset.y_train, 
                                enable_categorical=True,
                                weight=dataset.weight_train)
        dvalid_FS = xgb.DMatrix(dataset.X_valid[list(feat_imp.head(i).index)],dataset.y_valid, enable_categorical=True,)
    
        bst_FS = xgb.train(params, dtrain_FS,verbose_eval=False, num_boost_round=num_boost_round,
                                evals = [(dtrain_FS, "train"), (dvalid_FS, "valid")],
                                early_stopping_rounds=100,)
        preds_class_valid = bst_FS.predict(dvalid_FS)
        mcc = matthews_corrcoef(dataset.y_valid,preds_class_valid.argmax(1))
        acc = accuracy_score(dataset.y_valid,preds_class_valid.argmax(1))
        tot.append({'mcc':mcc,'acc':acc,'FS':i})
        
    tot = pd.DataFrame(tot)  
    FS = int(tot.loc[tot.acc.argmax()].FS) ## get best
    print(f'Best model with {FS} features, retraining....')
    
    dtrain_FS = xgb.DMatrix(dataset.X_train[list(feat_imp.head(FS).index)],dataset.y_train, enable_categorical=True, weight=dataset.weight_train)
    dvalid_FS = xgb.DMatrix(dataset.X_valid[list(feat_imp.head(FS).index)],dataset.y_valid,enable_categorical=True,  )
    bst_FS = xgb.train(params, dtrain_FS,verbose_eval=False, num_boost_round=num_boost_round,
                evals = [(dtrain_FS, "train"), (dvalid_FS, "valid")],
                early_stopping_rounds=100,)
    
        
    return tot,bst_FS,FS
