import numpy as np
import pandas as pd
from sklearn.covariance import EllipticEnvelope
from sklearn.svm import OneClassSVM
from sklearn.ensemble import IsolationForest

from pyod.utils.utility import standardizer #returns z-score


## Crossplot

def get_params(**kwargs):
    """
    Gets parameters for the outlier detection algorithms

    Returns:
        tuple: dictionaries with parameters per method
    """
    SVM_params = {
        'nu': kwargs.get('SVM_nu', 0.01)}
    IsoFor_params = {
        'n_estimators': kwargs.get('IsoFor_n_estimators', 50),
        'contamination': kwargs.get('IsoFor_contamination', 0.01),
        'random_state': kwargs.get('IsoFor_random_state', 0)
    }
    Abod_parameters = {
        'contamination': 0.1,
        'n_neighbors': 10
    }
    EliEnv_parameters = {
        'contamination': 0.05,
        'random_state': 0
    }
    return SVM_params, IsoFor_params, Abod_parameters, EliEnv_parameters

def find_crossplot_scores(df, x, y, algo='comb', **kwargs):
    """
    Find anomalies based on the crossplots of x vs. y

    Args:
        df_well (pd.DataFrame): dataframe with data from one well
        algo (str, optional): which algorithm to use. Defaults to 'comb'.

    Returns:
        list: list with indices of anomalous samples
    """
    algo = algo.lower()
    df  = df[(df[x]!=0) & (df[y]!=0)].copy()
    tmp = df[[x, y]].dropna().copy()

    anomal_idx = []
    method_scores = dict()
    return_scores = dict()
    SVM_params, IsoFor_params, Abod_parameters, EliEnv_parameters = get_params(**kwargs)
    
    #if there are not even 10 samples we return
    if len(tmp)<10:
        return [], {}, tmp.index

    if (algo=='comb') or (algo=='EliEnv'):
        try:
            elienv = EllipticEnvelope(
                contamination = EliEnv_parameters['contamination'],
                random_state = EliEnv_parameters['random_state']
            ).fit(tmp[[x, y]])    
            preds2 = elienv.predict(tmp[[x, y]])
            original_scores = elienv.score_samples(tmp[[x, y]])
        except: #FIXME: elliptic envelope throws covariance error when not so many samples.Quick fix for now.
            preds2 = np.zeros(len(tmp))
            original_scores = np.zeros(len(tmp))
        if algo=='EliEnv':
            preds = preds2
        method_scores['EliEnv'] = np.abs(standardizer(original_scores.reshape(-1,1)))

    if (algo=='comb') or (algo=='svm'):
        svm = OneClassSVM(
            kernel='rbf', 
            gamma='scale', 
            nu=SVM_params['nu']
        ).fit(tmp[[x, y]])
        preds3 = svm.predict(tmp[[x, y]])
        if algo=='svm':
            preds = preds3
        original_scores = svm.score_samples(tmp[[x, y]])
        method_scores['SVM'] = np.abs(standardizer(original_scores.reshape(-1,1)))

    if (algo=='comb') or (algo=='isofor'):
        ifo = IsolationForest(
            n_estimators=IsoFor_params['n_estimators'], 
            contamination=IsoFor_params['contamination'], 
            random_state=IsoFor_params['random_state']
        ).fit(tmp[[x, y]])
        preds4 = ifo.predict(tmp[[x, y]])
        if algo=='isofor':
            preds = preds4
        original_scores = ifo.score_samples(tmp[[x, y]])
        method_scores['IsoFor'] = standardizer(original_scores.reshape(-1,1))

    if algo=='comb':
        preds  = np.where(preds2==-1, 1, 0) +\
                np.where(preds3==-1, 1, 0) +\
                np.where(preds4==-1, 1, 0)
        tmp['pred'] = np.where(preds>1, 1, 0)
    else:
        tmp['pred'] = np.where(preds==-1, 1, 0)

    for k, v in method_scores.items():
        return_scores[k] = v[:,0]
    
    return_scores = pd.DataFrame(return_scores)
    return_scores['agg'] = return_scores.mean(axis=1)
    anomal_idx.append(list(tmp[tmp.pred==1].index))
    return sum(anomal_idx, []), return_scores, tmp.index

def train_crossplot(**kwargs):
    pass

def get_proba_predictions(X, y, bad_logs, model, cp_name):
    """
    Add to the dataframe the score results for each of the bad anomalies.

    Args:
        X (pd.DataFrame): dataframe for prediction on
        y (pd.DataFrame): datafrrame with columns to populate results
        bad_logs (pd.DataFrame): bad log samples dataframe
        models (dict): dictionary with models per cp
        cp_name (str): crossplot name

    Returns:
        y: pd.DataFrame populated with results
    """
    cp_preds = model.predict_proba(X)
    if cp_name == 'vp_den':
        y.loc[bad_logs.index, 'score_vpden_ac']  = cp_preds[:, 0]
        y.loc[bad_logs.index, 'score_vpden_den'] = cp_preds[:, 1]
    elif cp_name == 'vp_vs':
        y.loc[bad_logs.index, 'score_vpvs_ac']  = cp_preds[:, 0]
        y.loc[bad_logs.index, 'score_vpvs_acs'] = cp_preds[:, 1]
    elif cp_name == 'ai_vpvs':
        y.loc[bad_logs.index, 'score_aivpvs_ac']  = cp_preds[:, 0]
        y.loc[bad_logs.index, 'score_aivpvs_acs'] = cp_preds[:, 1]
        y.loc[bad_logs.index, 'score_aivpvs_den'] = cp_preds[:, 2]
    return y

def get_class_predictions(X, y, bad_logs, model, cp_name):
    """
    Add to the dataframe the class results for each of the bad anomalies.

    Args:
        X (pd.DataFrame): dataframe for prediction on
        y (pd.DataFrame): datafrrame with columns to populate results
        bad_logs (pd.DataFrame): bad log samples dataframe
        models (dict): dictionary with models per cp
        cp_name (str): crossplot name

    Returns:
        y: pd.DataFrame populated with results
    """
    cp_preds = model.predict(X)
    y['tmp'] = 0
    y.loc[bad_logs.index, 'tmp'] = cp_preds
    #Map predictions depending on which crossplot they come from
    if cp_name == 'vp_den':
        y.loc[y.tmp.isin([1,3]), 'flag_vpden_ac']  = 1
        y.loc[y.tmp.isin([2,3]), 'flag_vpden_den'] = 1        
    elif cp_name == 'vp_vs':
        y.loc[y.tmp.isin([1,3]), 'flag_vpvs_ac']  = 1
        y.loc[y.tmp.isin([2,3]), 'flag_vpvs_acs'] = 1 
    elif cp_name == 'ai_vpvs':
        y.loc[y.tmp.isin([4,5,6,7]), 'flag_aivpvs_den'] = 1
        y.loc[y.tmp.isin([1,3,5,7]), 'flag_aivpvs_ac']  = 1
        y.loc[y.tmp.isin([2,4,6,7]), 'flag_aivpvs_acs'] = 1
    return y

def flag_well_crossplots(X, x, y, algo, cp_name, **kwargs):
    """
    Returns anomalous indices of crossplots between x and y

    Args:
        X (pd.DataFrame): data of one well

    Returns:
        tuple: dataframe with scores and lists of anomalous indices
    """
    res = dict(zip(
        ['anomalies', 'scores', 'idx'],
        find_crossplot_scores(X, x, y, algo=algo, **kwargs)
    ))
    for comb_method, scs in res['scores'].items():
        method_col = '{}_{}'.format(comb_method, cp_name)
        X.loc[res['idx'], method_col] = scs.values
        X[method_col].fillna(0, inplace=True)
    return X, res['anomalies']

def flag_crossplots(X, x, y, algo, cp_name, **kwargs):
    """
    Returns anomalous indices of crossplots between x and y

    Args:
        X (pd.DataFrame): data of one well

    Returns:
        tuple: dataframe with scores and lists of anomalous indices
    """
    anomalies = []
    for lsu in X['GROUP'].unique():
        lsu_data = X[X['GROUP']==lsu]
        res = dict(zip(
            ['anomalies', 'scores', 'idx'],
            find_crossplot_scores(lsu_data, x, y, algo=algo, **kwargs)
        ))
        for comb_method, scs in res['scores'].items():
            method_col = '{}_{}'.format(comb_method, cp_name)
            X.loc[res['idx'], method_col] = scs.values
            X[method_col].fillna(0, inplace=True)
        anomalies.extend(res['anomalies'])
    return X, anomalies

def classify_badlog(y, cp_name, ds, key_wells, model):
    bad_logs = ds.df_original[ds.df_original['flag_crossplot_gen']!=0]
    # if no badlogs, dont do anything
    if len(bad_logs) == 0:
        return y
    df_preprocessed, _, feats = ds.preprocess(
        ds.df_original.loc[bad_logs.index],
        _normalize_curves={'key_wells':key_wells}
    )
    #scores
    y_scores = get_proba_predictions(df_preprocessed[feats], y, bad_logs, model, cp_name)
    y[y_scores.columns] = y_scores
    #flags
    y_flags = get_class_predictions(df_preprocessed[feats], y, bad_logs, model, cp_name)
    y[y_flags.columns] = y_flags
    return y

def flag_crossplot(X_pred, cp_names, datasets, models, key_wells, **kwargs):
    print('Method: crossplot - general...')
    outlier_methods = ['EliEnv', 'SVM', 'IsoFor'] #FIXME! The shouldn't preferably be hard-coded.
    for cp_name in cp_names:        
        for method in outlier_methods:
            X_pred['{}_{}'.format(method, cp_name)] = 0.

    #get anomalies per crossplot and their scores

    vp_den_anomalies, ai_vpvs_anomalies, vp_vs_anomalies = [], [], []
    
    X_pred, vp_den_anomalies  = flag_crossplots(X_pred, 'VP', 'DEN', 'comb', 'vp_den', **kwargs)

    X_pred, ai_vpvs_anomalies = flag_crossplots(X_pred, 'AI', 'VPVS', 'comb', 'ai_vpvs', **kwargs)

    X_pred, vp_vs_anomalies   = flag_crossplots(X_pred, 'VP', 'VS', 'comb', 'vp_vs', **kwargs)

    y = X_pred.copy()        
    y.loc[:, [
        'flag_vpden_ac', 'flag_vpvs_ac', 'flag_aivpvs_ac',
        'flag_vpvs_acs', 'flag_aivpvs_acs',
        'flag_vpden_den', 'flag_aivpvs_den'
    ]] = 0, 0, 0, 0, 0, 0, 0

    pred_gen_cols = ['flag_crossplot_gen', 'flag_vpden_gen', 'flag_aivpvs_gen', 'flag_vpvs_gen']
    y.loc[:, pred_gen_cols] = 0, 0, 0, 0        
    y.loc[vp_den_anomalies, ['flag_crossplot_gen', 'flag_vpden_gen']]   = 1
    y.loc[ai_vpvs_anomalies, ['flag_crossplot_gen', 'flag_aivpvs_gen']] = 1
    y.loc[vp_vs_anomalies, ['flag_crossplot_gen', 'flag_vpvs_gen']]     = 1

    #Add the general crossplot flags to X_pred(to use as features in the model)
    X_pred[pred_gen_cols] = y[pred_gen_cols]

    #Run the anomalous samples through the relevant trained model and get predictions
    #Preprocessing using key_wells will need the well_name column, but the value is unimportant
    X_pred['well_name'] = 'dummy'
    for cp_name in cp_names:
        print('Method: crossplot - {}...'.format(cp_name))
        ds = datasets[cp_name]
        ds.load_from_df(X_pred)
        y = classify_badlog(y, cp_name, ds, key_wells[cp_name], models[cp_name])
    y = y.fillna(0)

    return y
