from utils import retrive_data, split
from model import train, gain_accuracy_train
from sklearn.metrics import confusion_matrix,matthews_corrcoef,accuracy_score
import xgboost as xgb
import pandas as pd
import pickle
from omegaconf import DictConfig,OmegaConf
import hydra
import logging

import os




@hydra.main(config_name='config.yaml')
def main(conf: DictConfig) -> None:

    skiarea_test= conf.processing.skiarea_test 
    season_test_skiarea = conf.processing.season_test_skiarea
    season_test_year= conf.processing.season_test_year         
    weight_type = conf.processing.weight_type    
    reload_data = conf.processing.reload_data
    use_smote = conf.processing.use_smote         
    undersampling = conf.processing.undersampling 
    test_size = conf.processing.test_size
    use_small= conf.processing.use_small
    num_boost_round = conf.model.num_boost_round
    retrain_last_model = conf.model.retrain_last_model
    retrain = conf.model.retrain
    n_trials = conf.model.n_trials
    name = conf.model.name
    os.makedirs(name,exist_ok=True) 
    with open(os.path.join(name,"conf.yaml"),'w') as f:
        OmegaConf.save(conf, f)
    logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(os.path.join(name,"debug.log")),logging.StreamHandler() ])
    
    
    ## get the data
    labeled,labeled_small,to_remove,evacuations,encoders = retrive_data(reload_data=reload_data,threshold_under_represented=0.5,path='/home/agobbi/Projects/PID/datanalytics/PID/src')

    
    with open(os.path.join(name,'metadata.pkl'),'wb') as f:
        pickle.dump([to_remove,use_small,evacuations,encoders],f)

    #split the data

    dataset,dataset_test = split(labeled_small if use_small  else labeled  ,
                                skiarea_test= skiarea_test,
                                season_test_skiarea = season_test_skiarea,
                                season_test_year= season_test_year,
                                use_smote = use_smote,
                                undersampling = undersampling,
                                test_size = test_size,
                                weight_type = weight_type )
    #if you changed something you may want to retrain the model and save the best model
    if retrain:
        logging.info('OPTUNA hyperparameter tuning, please wait!')
        best_model,params_final,study = train(dataset,n_trials=n_trials,timeout=6000,num_boost_round=num_boost_round)
        feat_imp = pd.Series(best_model.get_fscore()).sort_values(ascending=False)

        with open(os.path.join(name,'best_params.pkl'),'wb') as f:
            pickle.dump([params_final,feat_imp,best_model,study],f)
    
    else:
        with open(os.path.join(name,'best_params.pkl'),'rb') as f:
            params_final,feat_imp,best_model,study = pickle.load(f)
            
    

    
    #for retriving prediction must convert to DMatrix type
    tmp_train = xgb.DMatrix(dataset.X_train[best_model.feature_names],dataset.y_train,enable_categorical=True)
    tmp_valid = xgb.DMatrix(dataset.X_valid[best_model.feature_names],dataset.y_valid,enable_categorical=True)

    ##get the scores
    preds_class_valid = best_model.predict(tmp_valid)
    preds_class_train= best_model.predict(tmp_train)
    logging.info('##################RESULT ON THE TRAIN SET#####################')
    logging.info(confusion_matrix(dataset.y_train,preds_class_train.argmax(1)))
    logging.info(f'MCC:{matthews_corrcoef(dataset.y_train,preds_class_train.argmax(1))}')
    logging.info(f'ACC:{accuracy_score(dataset.y_train,preds_class_train.argmax(1))}')
    logging.info('##################RESULT ON THE VALIDATION SET#####################')
    logging.info(confusion_matrix(dataset.y_valid,preds_class_valid.argmax(1)))
    logging.info(f'MCC:{matthews_corrcoef(dataset.y_valid,preds_class_valid.argmax(1))}')
    logging.info(f'ACC:{accuracy_score(dataset.y_valid,preds_class_valid.argmax(1))}')
    

        
    #now you can train the final model, for example using gain_accuracy_train for reducing the number of features used
    if retrain_last_model:
        tot,bst_FS,FS = gain_accuracy_train(dataset,feat_imp,num_boost_round=num_boost_round,params=params_final)
        with open(os.path.join(name,'best_params_and_final_model.pkl'),'wb') as f:
            pickle.dump([tot,bst_FS,FS],f)
            bst_FS.save_model(os.path.join(name,"model.json"))
    else:
        with open(os.path.join(name,'best_params_and_final_model.pkl'),'rb') as f:
            tot,bst_FS,FS = pickle.load(f)  
            bst_FS = xgb.Booster()
            bst_FS.load_model(os.path.join(name,"model.json"))


    ## save the model in json format, maybe it is better
    


    if dataset_test.X_test_area is not None:
        dtest_FS = xgb.DMatrix(dataset_test.X_test_area[bst_FS.feature_names],dataset_test.y_test_area,enable_categorical=True,)
        preds_class_test = bst_FS.predict(dtest_FS)
        mcc = matthews_corrcoef(dataset_test.y_test_area,preds_class_test.argmax(1))
        acc = accuracy_score(dataset_test.y_test_area,preds_class_test.argmax(1))
        cm = confusion_matrix(dataset_test.y_test_area,preds_class_test.argmax(1))
        logging.info(f'RESULT ON THE TEST SKI AREA {mcc=}, {acc=}, \n{cm=}')
        
    if dataset_test.X_test_season is not None:
        dtest_season_FS = xgb.DMatrix(dataset_test.X_test_season[bst_FS.feature_names],dataset_test.y_test_season,enable_categorical=True,)
        preds_class_test_season = bst_FS.predict(dtest_season_FS)
        mcc = matthews_corrcoef(dataset_test.y_test_season,preds_class_test_season.argmax(1))
        acc = accuracy_score(dataset_test.y_test_season,preds_class_test_season.argmax(1))
        cm = confusion_matrix(dataset_test.y_test_season,preds_class_test_season.argmax(1))

        logging.info(f'RESULT ON THE TEST SKI SEASON {mcc=}, {acc=}, {cm=}')
    

    
if __name__ == "__main__":
    main()
    