from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV, StratifiedKFold, RandomizedSearchCV
import time
from scipy.stats import randint
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.svm import SVC
import joblib
import sklearn.metrics as sm
from sklearn.metrics import confusion_matrix
import warnings
import os

NUM_SPLITS = 3  # Number of splits for cross-validation

def split_group(group):
    # Determine the split percentage based on group size
    if len(group) < 15:
        split_percentage = 1.0
    elif len(group) < 75:
        split_percentage = 0.95  # 5% split for groups with fewer than 75 rows
    elif len(group) < 100:
        split_percentage = 0.90  # 10% split for groups with fewer than 100 rows
    else:
        split_percentage = 0.80  # 80% split for larger groups

    # Calculate the training size
    train_size = int(len(group) * split_percentage)

    # Shuffle the group
    shuffled_group = group.sample(frac=1, random_state=42).reset_index(drop=True)

    # Assign splits based on the calculated training size
    shuffled_group.loc[:train_size - 1, 'split'] = 'train'
    shuffled_group.loc[train_size:, 'split'] = 'valid'

    return shuffled_group


def mark_train_valid(df, training_column):
    # Create a new column for the split
    df['split'] = np.nan

    # Function to split each group
    # Apply the function to each group
    df = df.groupby(training_column).apply(split_group).reset_index(drop=True)
    return df


def extract_train_features(df):
    df_train = df.loc[df["split"] == "train"]
    return df_train


def extract_validation_features(df):
    df_valid = df.loc[df["split"] == "valid"]
    return df_valid

def random_forest_training(x_train, y_train, log_filepath):

    t0 = time.time()
    logfile = open(log_filepath, "a")
    logfile.write("------\n")
    logfile.write("timestamp: {:s} \n".format(time.strftime("%Y-%m-%d %H:%M:%S", (time.localtime()))))
    logfile.write("name: {:s} \n".format(__file__))
    logfile.write("------ \n")
    logfile.write("trainset size: %d\n" % x_train.shape[0])
    logfile.write("------ \n")
    logfile.write("GridSearchCV...\n")


    # param_grid = {
    #     'n_estimators': [50, 100, 200],  # Number of trees in the forest
    #     'max_features': ['sqrt', 'log2', None],  # Max features to consider at each split
    #     'max_depth': [10, 20, 30, None],  # Max depth of trees
    #     'min_samples_split': [2, 5, 10],  # Minimum number of samples required to split a node
    #     'min_samples_leaf': [1, 2, 4],  # Minimum number of samples required at each leaf node
    #     'bootstrap': [True, False]  # Whether to bootstrap samples
    # }

    param_grid = {
        'rf__n_estimators': randint(50, 150),  # Narrower range for trees
        'rf__max_features': ['sqrt', 'log2'],  # Focus on common choices
        'rf__max_depth': [10, 20, None],  # Avoid redundant intermediate values
        'rf__min_samples_split': [2, 5],  # Use most commonly used values
        'rf__min_samples_leaf': [1, 2],  # Reduce options to most practical
        'rf__bootstrap': [True]  # Fix bootstrap to True for simplicity
    }


    cv = StratifiedKFold(n_splits=NUM_SPLITS, shuffle=True, random_state=42)

    pipeline = Pipeline([
        ('scaler', StandardScaler()),
        ('rf', RandomForestClassifier())
    ])

    paramter_search = RandomizedSearchCV(
        estimator=pipeline,
        param_distributions=param_grid,
        n_iter=25,  # Specify a smaller number of iterations
        scoring='accuracy',
        cv=cv,  # 5-fold cross-validation
        n_jobs=-1,
        random_state=42,
        verbose=2
    )

    # paramter_search = GridSearchCV(
    #     estimator=rf,
    #     param_grid=param_grid,
    #     scoring='accuracy',
    #     cv=cv,
    #     n_jobs=-1,
    #     verbose=2
    # )

    paramter_search.fit(x_train, y_train)

    t1 = time.time()
    seconds = t1 - t0
    minutes = seconds / 60
    hours = minutes / 60
    logfile.write("Processing time of SearchCV: %02d:%02d:%02d - %s [s]\n" % (hours, minutes % 60, seconds % 60, t1 - t0))

    # Extract the best hyperparameters
    best_params = paramter_search.best_params_
    print("Best Hyperparameters:", best_params)
    print("Best Cross-Validation Accuracy:", paramter_search.best_score_)
    logfile.write(f"Best parameter {best_params}\n")

    # Extract the best parameters for RandomForestClassifier
    # Remove the 'rf__' prefix from the parameter names to use them directly
    classifier_best_params = {key.split('__')[1]: value for key, value in best_params.items()}

    # Create a new pipeline with the best parameters
    best_pipeline = Pipeline([
        ('scaler', StandardScaler()),  # Keep the scaler as is
        ('rf', RandomForestClassifier(random_state=42, **classifier_best_params))  # Use the best parameters
    ])

    # Fit the new model on the training data
    best_pipeline.fit(x_train, y_train)
    logfile.close()

    return best_pipeline

def catboost_training(x_train, y_train):
    from catboost import CatBoostClassifier
    # Initialize the CatBoostClassifier
    catboost_model = CatBoostClassifier(iterations=100, 
                                        learning_rate=0.1, 
                                        depth=10, 
                                        loss_function='MultiClass', 
                                        verbose=False)

    # Train the model
    catboost_model.fit(x_train.astype(np.float32), y_train)
    
    return catboost_model

def svm_training(x_train, y_train, log_filepath, save_scaler_path=None, save_svc_path=None):
    t0 = time.time()
    logfile = open(log_filepath, "a")
    logfile.write("------\n")
    logfile.write("timestamp: {:s} \n".format(time.strftime("%Y-%m-%d %H:%M:%S", (time.localtime()))))
    logfile.write("name: {:s} \n".format(__file__))
    logfile.write("------ \n")
    logfile.write("trainset size: %d\n" % x_train.shape[0])
    logfile.write("------ \n")
    logfile.write("GridSearchCV...\n")

    param_grid = {
        'svc__C': [1e-1, 1, 5, 1e1, 1e2, 1e3, 5e3, 1e4, 5e4, 1e5],
        'svc__gamma': [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1],
        'svc__kernel': ['rbf']
    }

    cv = StratifiedKFold(n_splits=NUM_SPLITS, shuffle=True, random_state=42)

    pipeline = Pipeline([
        ('scaler', StandardScaler()),
        ('svc', SVC(probability=True))
    ])

    paramter_search = GridSearchCV(
        estimator=pipeline,
        param_grid=param_grid,
        scoring='accuracy',
        cv=cv,
        n_jobs=-1,
        verbose=2
    )

    paramter_search.fit(x_train, y_train)

    t1 = time.time()
    seconds = t1 - t0
    minutes = seconds / 60
    hours = minutes / 60
    logfile.write("Processing time of SearchCV: %02d:%02d:%02d - %s [s]\n" % (hours, minutes % 60, seconds % 60, t1 - t0))

    best_params = paramter_search.best_params_
    print("Best Hyperparameters:", best_params)
    print("Best Cross-Validation Accuracy:", paramter_search.best_score_)
    logfile.write(f"Best parameter {best_params}\n")

    # Remove the 'svc__' prefix from the parameter names
    classifier_best_params = {key.split('__')[1]: value for key, value in best_params.items()}

    # Save scaler and SVC separately if requested
    scaler = StandardScaler().fit(x_train)
    x_train_scaled = scaler.transform(x_train)
    svc = SVC(probability=True, random_state=42, **classifier_best_params)
    svc.fit(x_train_scaled, y_train)
    if save_scaler_path is not None:
        joblib.dump(scaler, save_scaler_path)
    if save_svc_path is not None:
        joblib.dump(svc, save_svc_path)

    # Also return a pipeline for convenience
    best_pipeline = Pipeline([
        ('scaler', scaler),
        ('svc', svc)
    ])

    logfile.close()
    return best_pipeline


def ComputeCmStatistics(classified, referenced, weight):
    """
    compute classification metrics (User's accuracy, Producer's accuracy, F-score, Omission Error, Commission Error, Overal Accuracy,Kappa Coefficient)

    :param classified: list with classified values
    :param referenced: list with reference values
    :param weight: list with areas (or None)
    :return: dictionary with classification metrics: 'UsersAccuracy':Users Accuracy, 'ProducersAccuracy':Producers Accuracy,
    'f_score':f-score,'kappa':Kappa Coefficient,'OveralAccuracy':Overal Accuracy,'OmissionError':Omission Error,
    'CommissionError':Commission Error
    """

    print('computing classification metrics..')
    # weighted version
    if not weight is None:
        print('computing classification metrics .. weighted version')
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            ProducersAccuracy = [round(elem, 2) for elem in list(
                sm.recall_score(referenced, classified, labels=None, pos_label=1, average=None, sample_weight=weight))]
            UsersAccuracy = [round(elem, 2) for elem in list(
                sm.precision_score(referenced, classified, labels=None, pos_label=1, average=None,
                                   sample_weight=weight))]
            OmissionError = []
            for itm in range(0, len(ProducersAccuracy)):
                OmissionError.append(round((1 - ProducersAccuracy[itm]), 2))
            CommissionError = []
            for itm in range(0, len(UsersAccuracy)):
                CommissionError.append(round((1 - UsersAccuracy[itm]), 2))
            f_score = [round(elem, 2) for elem in list(
                sm.f1_score(referenced, classified, labels=None, pos_label=1, average=None, sample_weight=weight))]
            OveralAccuracy = round(sm.accuracy_score(referenced, classified, normalize=True, sample_weight=weight), 2)
        kappa = "undefined"

    # non-weighted version
    else:
        print('computing classification metrics .. non-weighted version')
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            ProducersAccuracy = [round(elem, 2) for elem in list(
                sm.recall_score(referenced, classified, labels=None, pos_label=1, average=None, sample_weight=None))]
            UsersAccuracy = [round(elem, 2) for elem in list(
                sm.precision_score(referenced, classified, labels=None, pos_label=1, average=None, sample_weight=None))]
            OmissionError = []
            for itm in range(0, len(ProducersAccuracy)):
                OmissionError.append(round((1 - ProducersAccuracy[itm]), 2))
            CommissionError = []
            for itm in range(0, len(UsersAccuracy)):
                CommissionError.append(round((1 - UsersAccuracy[itm]), 2))
            f_score = [round(elem, 2) for elem in list(
                sm.f1_score(referenced, classified, labels=None, pos_label=1, average=None, sample_weight=None))]
            kappa = round(sm.cohen_kappa_score(referenced, classified, labels=None, weights=None), 2)
            OveralAccuracy = round(sm.accuracy_score(referenced, classified, normalize=True, sample_weight=None), 2)

    return {'UsersAccuracy': UsersAccuracy, 'ProducersAccuracy': ProducersAccuracy, 'f_score': f_score, 'kappa': kappa,
            'OveralAccuracy': OveralAccuracy, 'OmissionError': OmissionError, 'CommissionError': CommissionError}


def create_confusion_matrix(classified, referenced):
    # Compute the unique values for target_column and Final_Crop
    unique_target = np.unique(referenced)
    unique_pred = np.unique(classified)

    # Compute the confusion matrix
    cm_sklearn = confusion_matrix(referenced, classified, labels=unique_target)

    # Compute the unique values for target_column and Final_Crop
    unique_target = np.unique(referenced)

    # Convert the confusion matrix to a DataFrame for better formatting
    cm_df = pd.DataFrame(cm_sklearn,
                         index=unique_target,
                         columns=unique_target)

    # Fill missing columns for prediction classes not present in unique_target
    for col in unique_pred:
        if col not in cm_df.columns:
            cm_df[col] = 0

    # Rearrange columns to match the order of unique_pred
    cm_df = cm_df[unique_pred]

    return cm_df