
import pandas as pd
import psycopg2 as pg
import numpy as np
from sklearn.model_selection import train_test_split
import pickle
from dataclasses import dataclass
from typing import Union
import os
from imblearn.under_sampling import RandomUnderSampler,RandomOverSampler

##AUXILIARY CLASSES
@dataclass  
class Dataset:
    X_train:Union[pd.DataFrame,None]
    y_train:Union[pd.Series,None]
    X_valid:Union[pd.DataFrame,None]
    y_valid:Union[pd.Series,None]
    weight_train:Union[np.array, None]

@dataclass  
class Dataset_test:
    X_test_area:Union[pd.DataFrame,None]
    y_test_area:Union[pd.Series,None]
    X_test_season:Union[pd.DataFrame,None]
    y_test_season:Union[pd.Series,None]

def prepare_new_data(dataset:pd.DataFrame,to_remove:dict)->(pd.DataFrame,pd.DataFrame):
    """prepare new data for prediction. MUST BE SIMILAR TO retrive_data. Maybe it can use directly inside it...

    Args:
        dataset (pd.DataFrame): dataset to use as inference
        to_remove (dict): columns to aggregate

    Returns:
        two pandas dataframe, one the original the second with condensed classes. 

    """
    dataset_p = dataset.copy()
    dataset_p.drop(columns=['dateandtime','skiarea_id','day_of_year','minute_of_day','year'], inplace=True)
    
    ##evacuation_vehicles must be explicitated
    ev = set({})
    for i,row in dataset_p.iterrows():
        ev = ev.union(set(row.evacuation_vehicles))
    for c in ev:
        dataset_p[c] = False
    for i,row in dataset_p.iterrows():
        for c in row.evacuation_vehicles:
            dataset_p.loc[i,c] = True
    dataset_p.drop(columns=['town','province','evacuation_vehicles'],inplace=True)
    
    
    dataset_p['age'] =  dataset_p['age'].astype(np.float32).fillna(np.nan)
    
    dataset_p_small = dataset_p.copy()

    for c in to_remove.keys():
        for k in to_remove[c]:
            dataset_p_small.loc[dataset_p[c]==k,c] = 'other'
    for c in dataset_p.columns:
        if c not in ['age','season','skiarea_name']:
            dataset_p_small[c] =  dataset_p_small[c].fillna('None').astype('category')  
            dataset_p[c] =  dataset_p[c].fillna('None').astype('category')  
    dataset_p.dropna(inplace=True)
    dataset_p_small.dropna(inplace=True)

    return dataset_p,dataset_p_small



def retrive_data(reload_data:bool,threshold_under_represented:float,path:str)->(pd.DataFrame,pd.DataFrame):
    """Get data

    Args:
        reload_data (bool): if true, the procedure will downolad the data from the db
        threshold_under_represented (float): classes with few representants are condensed in the class `other`
        path (str): path in which saving the data

    Returns:
        two pandas dataframe, one the original the second with condensed classes and a dictionarly of condesed classes
    """
    if reload_data:
        engine = pg.connect("dbname='safeidx' user='fbk_mpba' host='172.104.247.67' port='5432' password='fbk2024$'")
        df = pd.read_sql('select * from fbk_export_20240212', con=engine) 
        with open(os.path.join(path,'data.pkl'),'wb') as f:
            pickle.dump(df,f)
    else:
        with open(os.path.join(path,'data.pkl'),'rb') as f:
            df = pickle.load(f)

    ## these columns can lead to overfit!
    df.drop(columns=['dateandtime','skiarea_id','day_of_year','minute_of_day','year'], inplace=True)
    
    ##evacuation_vehicles must be explicitated
    ev = set({})
    for i,row in df.iterrows():
        ev = ev.union(set(row.evacuation_vehicles))
    for c in ev:
        df[c] = False
    for i,row in df.iterrows():
        for c in row.evacuation_vehicles:
            df.loc[i,c] = True
    df.drop(columns=['town','province','evacuation_vehicles'],inplace=True)
    
    
    labeled = df[~pd.isna(df.india)].reset_index().drop(columns='index')
    labeled['age'] =  labeled['age'].astype(np.float32).fillna(np.nan)
    
    
    ## there are some under-represented categories, 
    ## maybe it is possible to obtain a more stable model removing such classes
    to_remove = {}
    for c in labeled.columns:
        if c not in ['india','age','season','skiarea_name']:
            labeled[c] = labeled[c].astype('str')  
            tmp = labeled.groupby(c)[c].count()
            tmp = 100*tmp/tmp.max()
            tmp = tmp[tmp<threshold_under_represented]
            if len(tmp)>0:
                for k in tmp.index:
                    if c not in to_remove.keys():
                        to_remove[c]= []
                    to_remove[c].append(k)
                    
    ## keep the datasets
    labeled_small = labeled.copy()
    for c in to_remove.keys():
        for k in to_remove[c]:
            labeled_small.loc[labeled_small[c]==k,c] = 'other'
    for c in labeled_small.columns:
        if c not in ['age','season','skiarea_name']:
            labeled_small[c] =  labeled_small[c].fillna('None').astype('category')  
            labeled[c] =  labeled[c].fillna('None').astype('category')
    labeled.dropna(inplace=True)
    labeled_small.dropna(inplace=True)
    
    #transform india into integers
    labeled.india = labeled.india.apply(lambda x: x.replace('i','')).astype(int)
    labeled_small.india = labeled_small.india.apply(lambda x: x.replace('i','')).astype(int)
    
    return labeled,labeled_small,to_remove
    


def split(labeled:pd.DataFrame,
          SKI_AREA_TEST: str = 'Klausberg',
          SEASON_TEST_SKIAREA:str = 'Kronplatz',
          SEASON_TEST_YEAR:int = 2023,
          use_smote:bool = False,
          undersampling:bool=False,
          test_size:float=0.33,
          weight_type:str = 'sqrt' )->(Dataset, Dataset_test):
    """Split  the dataset into train,validation test. From the initial dataset we remove a single skiarea (SKI_AREA_TEST)
    generating the first test set. Then we select a skieare and a starting season (SEASON_TEST_SKIAREA,SEASON_TEST_YEAR) 
    and generate the seconda test set. The rest of the data are splitted 66-33 stratified on the target column (india). 
    It is possible to specify the weight of eact sample. There are two strategies implemented: using the sum or the square root 
    of the sum. This is used for mitigating the class umbalance. Another alternative is to use an oversampling procedure (use_smote)

    Args:
        labeled (pd.DataFrame): dataset
        SKI_AREA_TEST (str, optional): skiarea to remove from the train and use in test. Defaults to 'Klausberg'.
        SEASON_TEST_SKIAREA (str, optional): skiarea to remove from the dataset if the season is greater than SEASON_TEST_YEAR. Defaults to 'Kronplatz'.
        SEASON_TEST_YEAR (int, optional): see SEASON_TEST_SKIAREA . Defaults to 2023.
        use_smote (bool, optional): use oversampling for class umbalance. Defaults to False.
        undersampling (bool, optional): use undersampling for class umbalance. Defaults to False.
        test_size (float, optional): percentage of dataset to use as validation. Defaults to 0.33.
        weight_type (str, optional): routine for weighting the error on the samples. Defaults to 'sqrt'.


    Returns:
        trainin-validation dataset and test dataset
    """
    labeled_tmp = labeled.copy()
    ##remove from dataset the corresponding test rows
    
    if SKI_AREA_TEST is not None:
        test_area = labeled[labeled.skiarea_name==SKI_AREA_TEST]
        labeled_tmp = labeled_tmp[labeled_tmp.skiarea_name!=SKI_AREA_TEST]
    else:
        test_area = None
        
    if SEASON_TEST_SKIAREA is not None and SEASON_TEST_YEAR is not None:
        test_area_season = labeled[(labeled.skiarea_name==SEASON_TEST_SKIAREA)&(labeled.season>=SEASON_TEST_YEAR)]
        labeled_tmp = labeled_tmp[(labeled_tmp.skiarea_name!=SEASON_TEST_SKIAREA)|(labeled_tmp.season<SEASON_TEST_YEAR) ]
    else:
        test_area_season = None

    
    
    

        
    
    X_train, X_valid, y_train, y_valid = train_test_split( labeled_tmp.drop(columns=['india','season','skiarea_name']),
                                                        labeled_tmp.india, test_size=test_size, random_state=0,stratify=labeled_tmp.india)
    
    if use_smote:   
        sm = RandomOverSampler()
        X_train,y_train = sm.fit_resample(X_train,y_train)
    if undersampling:
        sm = RandomUnderSampler(sampling_strategy='majority')
        X_train,y_train = sm.fit_resample(X_train,y_train)

    ##computed the weights for unbalanced dataset

    w = pd.DataFrame(np.unique(y_train,return_counts=True)).T
    w.columns = ['class','p']
    
    ##when computing the error, these are the weights used for each class: you can punish more errpr on most severe clases
    if weight_type == 'sqrt':
        w.p = np.sqrt(w.p.sum())/w.p
        print(w)
    elif weight_type == 'sum':
        w.p = w.p.sum()/w.p/w.shape[0]
        print(w)
    else:
        print(f'{weight_type=} not implemented please use a valid one: sqrt or sum, I will set all the weights to 0')
        w.p = 1
        
    if use_smote is False and undersampling is False:
        weight_train = pd.merge(pd.DataFrame({'class':y_train}),w).p.values
    else:
        w.p = 1
        weight_train = pd.merge(pd.DataFrame({'class':y_train}),w).p.values
    
        
    dataset = Dataset(X_train, y_train, X_valid, y_valid,weight_train)
    dataset_test = Dataset_test(test_area,test_area.india if test_area is not None else None,test_area_season,test_area_season.india if test_area_season is not None else None)

    return dataset,dataset_test

