
from autopilotml.data_loading import read_csv, read_excel, read_sqlite, read_postgres, read_mysql, read_mongo
from autopilotml.preprocessing import drop_missing_rows, drop_missing_columns, imputation, inter_quartile_range, remove_outliers_zscore
from autopilotml.transformation import ordinal_transform, onehot_transform, label_transform
from autopilotml.scaling import standard_scale, target_scale, robust_scale, maxabs_scale, power_transform, quantile_transform
from autopilotml.feature_selection import rfe, rfecv
from autopilotml.model_building import model_fitting
from autopilotml.model_tuning import tuner
from autopilotml.tracking import model_tracking

from sklearn.metrics import mean_squared_error, accuracy_score
import optuna
import mlflow
from mlflow import MlflowClient
from datetime import datetime
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
import warnings
warnings.filterwarnings("ignore")



#Stage 1: Data Loading

def load_data(path, csv = True, excel = False, **kwargs):
    """ This function is used to load the data from the source - locally stored files with extension csv or xlsx.

    Args:
        path (str): The path to the source file.

        csv (bool): If True, the source file is a CSV file.

        excel (bool): If True, the source file is an Excel file.

        **kwargs: Keyword arguments to pass to pandas.read_csv or pandas.read_excel.

    Returns:
        pandas.DataFrame: The DataFrame containing the data.
    """
    # Load the data from the csv source
    if excel:
        df = read_excel(path, **kwargs)
    elif csv:
        df = read_csv(path, **kwargs)

    return df

def load_database(database_type, query, sqlite_db_path= None, host= None, port= None, database_name= None, 
                  username= None, password=None, collection_name= None, **kwargs):
    """ This function is used to load the data from the database.

    Args:
        database (str): The type of database. For instance 'sqlite', 'mysql', 'postgres', 'mongo'.

        query (str): The query to execute.

        sqlite_db_path (str): The path to the SQLite database.

        The following arguments are used for MySQL, PostgreSQL and MongoDB databases:

        host (str): The host name.

        port (int): The port number.

        database_name (str): The database name.

        username (str): The username if exists. by default None.

        password (str): The password if exists. by default None.

        collection_name (str): The collection name for mongodb database. by default None.

        **kwargs: Keyword arguments to pass to sqlite3.connect, psycopg2.connect, mysql.connector.connect, pymongo.MongoClient.

    Returns:
        pandas.DataFrame: The DataFrame containing the data.
    """
    try:
        if database_type == 'sqlite':
            df = read_sqlite(query, sqlite_db_path)
        elif database_type == 'postgres':
            df = read_postgres(query, host, port, database_name, username, password)
        elif database_type == 'mysql':
            df = read_mysql(query, host, port, database_name, username, password)
        elif database_type == 'mongo':
            df = read_mongo(query, host, port, database_name, username, password, collection_name)
        else:
            raise ValueError('The database type is not supported.')
        
        return df
        
    except ValueError as e:
        print(e)
    except Exception as e:
        print(e)


#Stage 2: Data Preprocessing

def preprocessing(dataframe, label_column,
                missing= {'type': 'drop',
                        'drop_columns': True, 
                        'threshold': 0.25, 
                        'strategy_numerical': 'knn',
                        'strategy_categorical': 'most_frequent',
                        'fill_value': None}, 
                outlier= {'method': 'None',
                        'zscore_threshold': 3,
                        'iqr_threshold': 1.5,
                        'Lc': 0.25, 
                        'Uc': 0.75,
                        'cap': False},
                **kwargs):
    
    """ This function is used to preprocess the data. The preprocessing steps include missing value imputation and 
    outlier removal.

    Args:
        dataframe (pandas.DataFrame): The DataFrame containing the data.

        label_column (str): The name of the label column (target feature).

        missing (dict): The dictionary containing the imputation parameters. If you want to change any parameters in the dictionary, you have to pass the entire dictionary with the new values or default values.

            The following parameters are supported:
            
            type (str): The type of imputation. Can be one of "drop", "impute".
                drop: The rows with missing values will be dropped.
                impute: The missing values will be imputed using the imputation strategy.

            drop_columns (bool): If True, the columns with missing values greater than the threshold will be dropped.

            threshold (float): The threshold to use for dropping columns with missing values. Must be between 0 and 1.

            strategy_numerical (str): The imputation strategy to use for numerical columns. Can be one of "mean", "median", "most_frequent", "knn" or "constant".

            strategy_categorical (str): The imputation strategy to use for categorical columns. Can be one of "most_frequent", or "constant".
            
            fill_value (str/int/float): The (str) value to use for imputing missing values when strategy is "constant for object data types, For numerical it should be a int/float".

        outlier (dict): The dictionary containing the outlier parameters. If you want to change any parameters in the dictionary, you have to pass the entire dictionary with the new values or default values.

            The following parameters are supported:

            method (str): The method to use for removing outliers. Can be one of "zscore", "iqr", "None".
                None = No outlier removal.

            iqr_threshold (float): The threshold to use for removing outliers. Must be greater than 0.

            Lc (float): The lower cut-off percentile of iqr value to use for removing outliers. Must be between 0 and 1. By deafult 0.25.

            Uc (float): The upper cut-off percentile of iqr value to use for removing outliers. Must be between 0 and 1. By deafult 0.75.

            cap (bool): If True, the outliers will be capped to the upper and lower bounds. If False, the outliers will be removed.

        **kwargs: Keyword arguments to pass to sklearn KNNImputer.

    Returns:
        pandas.DataFrame: The DataFrame containing the preprocessed data.
    """
    
    data = dataframe.copy()

    if missing['drop_columns']:
        data = drop_missing_columns(data, threshold=missing['threshold'])

    if data[label_column].isnull().sum() > 0:
        print('The Label column has missing values. so dropping the rows with missing values in label column')
        data = drop_missing_rows(data,subset=[label_column])


    # Separate columns into numerical and categorical
    numerical_columns = data.select_dtypes(exclude='object').columns
    categorical_columns = data.select_dtypes(include='object').columns

    # Remove label column from numerical and categorical columns
    if data[label_column].dtype == 'object':
        categorical_columns = categorical_columns.drop(label_column)
    else:
        numerical_columns = numerical_columns.drop(label_column)

    if missing['type'] == 'drop':
        data = drop_missing_rows(data)

    elif missing['type'] == 'impute':
        data =  imputation(data, strategy_numerical=missing['strategy_numerical'], 
                            strategy_categorical=missing['strategy_categorical'],
                            fill_value=missing['fill_value'],
                            numerical_columns=numerical_columns,
                            categorical_columns=categorical_columns, 
                            **kwargs)
    if outlier['method'] == 'None':
        pass
    elif outlier['method'] == 'zscore':
        data = remove_outliers_zscore(data, numerical_columns=numerical_columns, zscore_threshold=outlier['zscore_threshold'])
    elif outlier['method'] == 'iqr':
        data = inter_quartile_range(data, numerical_columns=numerical_columns, iqr_threshold=outlier['iqr_threshold'], Lc=outlier['Lc'], Uc=outlier['Uc'], cap=outlier['cap'])
        
    return data

#Stage 3: Data Transformation

def transformation(dataframe, label_column, type = 'ordinal',target_transform = False, cardinality = True, Cardinality_threshold = 0.3):
    """ This function is used to transform the feature and target data. The target column will be transformed only if the target_transform is True. Only applicable for Classification usecase.
    If target_transform is True, the function will return the 3 objects, dataframe, encoder and target_encoder. If target_transform is False, the function will return the 2 objects, dataframe and encoder.

    Args:
        dataframe (pandas.DataFrame): The DataFrame containing the data.

        label_column (str): The name of the label column (target column).

        type (str): The type of transformation. Can be one of "ordinal", "onehot", "None". By default "ordinal".

        target_transform (bool): If True, the label column will be transformed.

        cardinality (bool): If True, the categorical columns with cardinality greater than the threshold will be dropped.

        Cardinality_threshold (float): The threshold to use for dropping columns with cardinality (how many unique values). Must be between 1 to any positive value. (e.g) If the thresold is 30. then any categorical column has more than 30 unique values will be dropped.

    Returns:
        pandas.DataFrame: The DataFrame containing the transformed data.
        encoder: The encoder object used for transforming the features.
        target_encoder: The encoder object used for transforming the target column. Applicable only for Classification usecase.
    """
    categorical_columns = dataframe.select_dtypes(include=['object']).columns
    categorical_columns = categorical_columns.drop([label_column]) if label_column in categorical_columns else categorical_columns


    if cardinality:
        drop_columns = [x for x in categorical_columns if dataframe[x].nunique() > Cardinality_threshold]
        # drop_columns = [x for x in drop_columns if x not in label_column or x in categorical_columns]
        print('List of columns dropped due to high cardinality: {}'.format(drop_columns))
        dataframe.drop(drop_columns, axis=1, inplace=True)
        categorical_columns = categorical_columns.drop(drop_columns)

    if type == 'ordinal':
        dataframe, encoder = ordinal_transform(dataframe, categorical_columns)
    elif type == 'onehot':
        dataframe, encoder = onehot_transform(dataframe, categorical_columns)
    elif type == 'None':
        print('No transformation applied on features')
    else:
        raise ValueError('The transformation type is not supported.')

    if target_transform:
        dataframe, target_encoder = label_transform(dataframe, label_column)

        return dataframe, encoder, target_encoder
    
    return dataframe, encoder

#Stage 4: Data Scaling

def scaling(dataframe, label_column, type = 'standard', target_scaling = False):
    """ This function is used to scale the feature and target data. when target_scaling is True, the function will return the 3 objects, dataframe, scaler and target_scaler. 
    If target_scaling is False, the function will return the 2 objects, dataframe and scaler.
    Target scaling is applicable only for Regression usecase.

    Args:
        dataframe (pandas.DataFrame): The DataFrame containing the data.

        label_column (str): The name of the label column.

        type (str): The type of scaling. Can be one of "standard", "robust", "maxabs", "power", "quantile".

        target_scaling (bool): If True, the label column will be scaled using Sklearn MinMax Scaler. Applicable only for Regression usecase.

    Returns:
        pandas.DataFrame: The DataFrame containing the scaled data.
        scaler: The scaler object used for scaling the features.
        target_scaler: The scaler object used for scaling the target column. Applicable only for Regression usecase.
    """
    numerical_columns = dataframe.select_dtypes(exclude='object').columns
    numerical_columns = numerical_columns.drop([label_column]) if label_column in numerical_columns else numerical_columns

    if type == 'standard':
        dataframe, scaler = standard_scale(dataframe, numerical_columns)
    elif type == 'robust':
        dataframe, scaler = robust_scale(dataframe, numerical_columns)
    elif type == 'maxabs':
        dataframe, scaler = maxabs_scale(dataframe, numerical_columns)
    elif type == 'power':
        dataframe, scaler = power_transform(dataframe, numerical_columns)
    elif type == 'quantile':
        dataframe, scaler = quantile_transform(dataframe, numerical_columns)
    elif type == 'None':
        print('No scaling applied on features')
    else:
        raise ValueError('The scaling type is not supported.')
    
    if target_scaling:
        dataframe, target_scaler = target_scale(dataframe, label_column)

        return dataframe, scaler, target_scaler
    else:
        return dataframe, scaler

#Stage 5: Feature Engineering


#Stage 6: Feature Selection
def feature_selection(dataframe, label_column, estimator, type='rfecv', max_features=10, min_features=2, scoring= 'accuracy', 
                    cv=15, n_jobs= -1, **kwargs):
    """ This function is used to select the top n features from the dataframe.

    Args:
        dataframe (pandas.DataFrame): The DataFrame containing the data.

        label_column (str): The name of the label column.

        estimator (str): The estimator to use for feature selection. Can be one of "LinearRegression", "BayesianRidge", 
            "RandomForestRegressor", "XGBRegressor", "GradientBoostingRegressor", "DecisionTreeRegressor", 
            "LogisticRegression", "RandomForestClassifier", "XGBClassifier",
            "GradientBoostingClassifier", "DecisionTreeClassifier".

        type (str): The type of feature selection. Can be one of "rfe", "rfecv".

        max_features (int): The maximum number of features to select.

        min_features (int): The minimum number of features to select.

        scoring (str): The scoring function to use for feature selection. Can be one of "accuracy", "f1",
            "precision", "recall", "roc_auc", "neg_mean_squared_error", "neg_mean_absolute_error", "neg_median_absolute_error", "r2".

        cv (int): The number of folds to use for cross-validation.By  default 15.

        n_jobs (int): The number of jobs to run in parallel.
        
        **kwargs: Keyword arguments to pass to sklearn.feature_selection.RFE or sklearn.feature_selection.RFECV.

    Returns:
        pandas.DataFrame: The DataFrame containing the selected features.
        selector: The selector object used for selecting the features.
    """
    if type == 'rfe':
        df, selector = rfe(dataframe, label_column, estimator, n_features_to_select=max_features, step=1, verbose=0, **kwargs)
    elif type == 'rfecv':
        df, selector = rfecv(dataframe, label_column, estimator, scoring, cv, min_features_to_select = min_features, n_jobs=n_jobs, **kwargs)
    elif type == 'None':
        print('No feature selection applied')
        return None
    else:
        raise ValueError('The feature selection type is not supported.')
    
    return df, selector

#Stage 7: Model Training

def training(dataframe, label_column, model_name, problem_type, target_scaler=None, test_split =0.15, hypertune=True, n_epochs = 100):
    """ This function is used to train the model.

    Args:
        dataframe (pandas.DataFrame): The DataFrame containing the data.

        label_column (str): The name of the label column.

        model_name (str): The name of the model. Can be one of "LinearRegression", "BayesianRidge", "RandomForestRegressor", "XGBRegressor",
                            "GradientBoostingRegressor", "SVR", "DecisionTreeRegressor", "KNeighborsRegressor", "LogisticRegression", 
                            "RandomForestClassifier", "XGBClassifier", "GradientBoostingClassifier", "SVC", "DecisionTreeClassifier", 
                            "KNeighborsClassifier", "GaussianNB".

        problem_type (str): The type of problem. Can be one of "Regression", "Classification".

        target_scaler : The scaler object name used for scaling the label column during scaling phase. Applicable only for Regression usecase. 

        test_split (float): The percentage of data to use for testing. it should be 0 to 1.

        hypertune (bool): If True, the model will be hypertuned using Optuna.

        n_epochs (int): The number of epochs to use for hypertuning.

    Returns:
        model: The trained model.
    """
    # Split the dataset into train and test
    x_train, x_test, y_train, y_test = train_test_split(dataframe.drop(label_column, axis=1), dataframe[label_column], 
                                                        test_size=test_split, random_state=42)

    if "XGB" in model_name:
        mlflow.xgboost.autolog(log_input_examples=True, log_model_signatures=True, log_models=True, 
        disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False
        )
    else:
        mlflow.sklearn.autolog(log_input_examples=True, log_model_signatures=True, log_models=True, 
        disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False,log_post_training_metrics=True
        )


    if hypertune:

        def objective(trial):

            tuner_params = tuner(model_name, trial)

            model, y_pred = model_fitting(x_train, x_test, y_train, y_test, model_name= model_name, params= tuner_params)
            if problem_type == 'Regression':
                if target_scaler is not None:
                    y_pred = target_scaler.inverse_transform(y_pred.reshape(-1,1))
                    actual = target_scaler.inverse_transform(np.array(y_test).reshape(-1,1))
                accuracy = round(mean_squared_error(actual, y_pred, squared=False),4)
            else:
                accuracy = round(accuracy_score(y_test, y_pred),4)

            return accuracy

        study = optuna.create_study(direction='minimize' if problem_type == 'Regression' else 'maximize')
        study.optimize(objective, n_trials=n_epochs)

        best_params = study.best_params

        print('Best Score: {}'.format(study.best_value))
        print('Best Parameters: {}'.format(study.best_params))

        model, y_pred = model_fitting(x_train, x_test, y_train, y_test, model_name= model_name, params= best_params)


    else:
        model, y_pred = model_fitting(x_train, x_test, y_train, y_test, model_name= model_name)

        if problem_type == 'Regression':
            if target_scaler is not None:
                actual = target_scaler.inverse_transform(np.rray(y_test).reshape(-1,1))
                y_pred = target_scaler.inverse_transform(y_pred.reshape(-1,1))
            accuracy = round(mean_squared_error(actual, y_pred, squared=False),4)
        else:
            accuracy = round(accuracy_score(y_test, y_pred),4)
        
        print('Accuracy: {}'.format(accuracy))

    return model

