from sklearn.model_selection import StratifiedKFold, GridSearchCV
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from sklearn.naive_bayes import GaussianNB, BernoulliNB, CategoricalNB
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from catboost import CatBoostClassifier
from xgboost import XGBClassifier
from wittgenstein import RIPPER
import pandas as pd
from typing import Literal
from sklearn.base import ClassifierMixin, BaseEstimator
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder, LabelEncoder, FunctionTransformer, StandardScaler
import math
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from onnxmltools import convert_xgboost


#=========================== TransformTargetClassifier ===========================#

class TransformedTargetClassifier(ClassifierMixin, BaseEstimator) :

    """
    A wrapper class used for encoding and inversely transforming predictions to the original label
    """

    def __init__(self, classifier, transformer) :
        self.classifier = classifier
        self.transformer = transformer

    def fit(self, x, y, **kwargs) :
        y = self.transformer.fit_transform(y)
        return self.classifier.fit(x, y, **kwargs)
    
    def predict(self, x_test) :
        if not isinstance(x_test, pd.DataFrame):
            if hasattr(self.classifier, "feature_names_in_"):
                cols = self.classifier.feature_names_in_
            elif hasattr(self.classifier, "feature_names_"):
                cols = self.classifier.feature_names_
            else:
                cols = None
            x_test = pd.DataFrame(x_test, columns=cols)
        output = self.classifier.predict(x_test)
        return self.transformer.inverse_transform(output)
    
    def predict_proba(self, x):
        proba = self.classifier.predict_proba(x)
        return proba / proba.sum(axis=1, keepdims=True)
    
    def __getattr__(self, name):
        return getattr(self.classifier, name)
    


#=========================== AutofitWrapper ===========================#

class _AutofitWrapper :

    def __init__(self, m, s, clf, x):
        self.model = m
        self.score = s
        self.classifier = clf
        current_wrapper = self.model
        self.convertible_model = self.model
        if isinstance(self.convertible_model, VotingClassifier):
            raise NotImplementedError("ONNX conversion directly not supported. Export the base estimators separately instead!")
        if isinstance(current_wrapper, TransformedTargetClassifier):
            current_wrapper = current_wrapper.classifier
            self.convertible_model = self.convertible_model.classifier
        if isinstance(current_wrapper, Pipeline):
            if "preprocess" in current_wrapper.named_steps:
                current_wrapper = current_wrapper.named_steps["preprocess"]
            if "model" in self.convertible_model.named_steps:
                self.convertible_model = self.convertible_model.named_steps["model"]
        if not isinstance(current_wrapper, ColumnTransformer):
            self.n_features = x.shape[1]
            self.preprocessedX = x
        else:
            self.preprocessedX = current_wrapper.transform(x)
            self.n_features = len(current_wrapper.get_feature_names_out())
        self.initial_type = [("input", FloatTensorType([None, self.n_features]))]
        
    def export_to_onnx(self, disable_zipmap:bool=True):

        """
        Dump the model as 'model.onnx' in your current working directory (Input name: 'input')
        """

        if not isinstance(self.convertible_model, (XGBClassifier, CatBoostClassifier, RIPPER)):
            onnx_model = convert_sklearn(self.convertible_model, initial_types=self.initial_type, options={id(self.convertible_model): {"zipmap":disable_zipmap}})
        elif isinstance(self.convertible_model, XGBClassifier):
            onnx_model = convert_xgboost(self.convertible_model, initial_types=self.initial_type, options={id(self.convertible_model): {"zipmap":disable_zipmap}})
        elif isinstance(self.convertible_model, CatBoostClassifier):
            onnx_model = None
        else:
            raise NotImplementedError("ONNX conversion not supported")
            
        if onnx_model:
            with open("model.onnx", "wb") as f:
                f.write(onnx_model.SerializeToString())
        else:
            self.convertible_model.save_model("model.onnx", format="onnx")



#=========================== autofit_classification_model ===========================#

def autofit_classification_model(x:pd.DataFrame, y:pd.DataFrame, x_type:Literal["numeric", "categorical", "mixed"], n_splits:int, cat_features:list[str]=[], whitelisted_algorithms:list[Literal["LogisticRegression", "DecisionTree", "RandomForest", "GaussianNB", "BernoulliNB", "CategoricalNB", "CatBoost", "XGBoost", "Ripper", "SVC", "KNN"]]|Literal["auto"]="auto", enable_votingClassifier:bool=True, random_state:int|None=42, verbosity:bool=True) -> object :
    
    """
    Automatically trains classification models with the best algorithm and hyperparameters based on ROC-AUC score
    """

    if (x_type=="numeric"):
        assert len(cat_features) == 0, "cat_features should be None for numeric x_type"
    else:
        assert len(cat_features) > 0, "cat_features cannot be None or Empty while x_type is not numeric"

    if (n_splits<2 or type(n_splits) is not int):
        raise ValueError ('n_splits must be integer and greater than 1')
    else:
        cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)

    if enable_votingClassifier:
        assert len(whitelisted_algorithms) > 2, "minimum 3 algorithms required to enable votingClassifier"

    score='roc_auc_ovr_weighted'
    numeric_features = []

    for i in x.columns:
        if i not in cat_features:
            numeric_features.append(i)

    def get_target_transformer(y):
        if y.dtype.kind in 'iu':
            return FunctionTransformer(validate=False)
        else:
            return LabelEncoder()

    def lr_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(drop="first", handle_unknown='ignore', sparse_output=False), cat_features), ('num', StandardScaler(), numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', LogisticRegression())
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
                'classifier__model__C':[0.01, 0.1, 1, 10, 100],
                'classifier__model__class_weight':[None, 'balanced'],
                'classifier__model__solver':['lbfgs','newton-cg', 'saga'],
                'classifier__model__max_iter': [100000]
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying LogisticRegression...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "LogisticRegression"
        return grid.best_score_

    def dt_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(dtype=int, handle_unknown='ignore', sparse_output=False), cat_features), ('num', 'passthrough', numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', DecisionTreeClassifier())
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
                'classifier__model__ccp_alpha':[0.0, 0.001],
                'classifier__model__class_weight':[None, 'balanced']
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying DecisionTree...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "DecisionTreeClassifier" 
        return grid.best_score_

    def rf_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(dtype=int, handle_unknown='ignore', sparse_output=False), cat_features), ('num', 'passthrough', numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', RandomForestClassifier())
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
                'classifier__model__n_estimators':[int(len(y)/5), int(len(y)/3.33), int(len(y)/2.5)],
                'classifier__model__ccp_alpha':[0.0, 0.001],
                'classifier__model__class_weight':[None, 'balanced'],
                'classifier__model__max_features':[None]
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying RandomForest...")
        grid.fit(x, y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "RandomForestClassifier"
        return grid.best_score_

    def GaussianNB_model(get_model=False) :
        clf = TransformedTargetClassifier(classifier=GaussianNB(), transformer=get_target_transformer(y))
        grid = GridSearchCV(clf, param_grid={}, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying GaussianNB...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "GaussianNB"
        return grid.best_score_

    def BernoulliNB_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(dtype=int, handle_unknown='ignore', sparse_output=False), cat_features), ('num', 'passthrough', numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', BernoulliNB())
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {'classifier__model__alpha':[0.001,0.01,0.1,1]}
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying BernouliNB...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "BernoulliNB"
        return grid.best_score_
    
    def CategoricalNB_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(dtype=int, handle_unknown='ignore', sparse_output=False), cat_features), ('num', 'passthrough', numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', CategoricalNB())
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
            'classifier__model__alpha':[0.001,0.01,0.1,1],
            'classifier__model__fit_prior':[True, False]
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying CategoricalNB...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "CategoricalNB"
        return grid.best_score_

    def CatBoost_model(get_model=False) :
        loss = "Logloss" if y.nunique() == 2 else "MultiClassOneVsAll"
        eval_metric = "AUC" if y.nunique() == 2 else "MultiClassOneVsAll"
        cat = CatBoostClassifier(
            loss_function=loss,
            cat_features=cat_features,
            verbose=False,
            random_seed=random_state,
            eval_metric=eval_metric,
            early_stopping_rounds=50,
            depth=8,
            allow_writing_files=False,
            bagging_temperature=0
        )
        clf = TransformedTargetClassifier(classifier=cat, transformer=FunctionTransformer(validate=False))
        param = {
            'classifier__learning_rate':[0.01,0.03,0.06,0.1],
            'classifier__random_strength':[0,1],
            'classifier__l2_leaf_reg':[1,3,5,8]
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying CatBoost...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "CatBoost"
        return grid.best_score_
    
    def XGBoost_model(get_model=False) :
        obj = "binary:logistic" if y.nunique() == 2 else "multi:softprob"
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), cat_features), ('num', 'passthrough', numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', XGBClassifier(max_depth=8, random_state=random_state, objective=obj, num_class=y.nunique() if y.nunique() > 2 else None))
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
            'classifier__model__n_estimators':[1000],
            'classifier__model__learning_rate':[0.01,0.03,0.06,0.1],
            'classifier__model__min_child_weight':[1,5,10],
            'classifier__model__gamma':[0,0.5],
            'classifier__model__reg_alpha':[0,0.01,0.1,1],
            'classifier__model__reg_lambda':[0,0.1,1,10]
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying XGBoost...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "XGBoost"
        return grid.best_score_
    
    def Ripper_model(get_model=False) :
        if (y.nunique() > 2):
            return -1000
        clf = TransformedTargetClassifier(classifier=RIPPER(random_state=random_state, k=10, dl_allowance=math.log2(len(y)*len(x.columns))), transformer=get_target_transformer(y))
        param = {
            'classifier__n_discretize_bins':[10, 20],
            'classifier__prune_size':[0.1,0.25,0.5,0.75],
        }
        grid = GridSearchCV(clf, param_grid=param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying Ripper...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "Ripper"
        return grid.best_score_

    def SVC_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), cat_features), ('num', StandardScaler(), numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', SVC(random_state=random_state, probability=True))
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
            'classifier__model__C':[0.01, 0.1, 1, 10, 100],
            'classifier__model__kernel':['linear', 'rbf', 'poly', 'sigmoid'],
            'classifier__model__coef0':[-1, -0.5, 0, 0.5, 1],
            'classifier__model__gamma':['scale', 'auto'],
            'classifier__model__class_weight':[None, 'balanced']
        }
        grid = GridSearchCV(clf, param_grid=param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying SVC...")
        grid.fit(x,y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "SVC"
        return grid.best_score_
    
    def KNN_model(get_model=False) :
        preprocessor = ColumnTransformer(transformers=[('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), cat_features), ('num', StandardScaler(), numeric_features)])
        preprocessor.set_output(transform="pandas")
        pipeline = Pipeline([
            ('preprocess', preprocessor), 
            ('model', KNeighborsClassifier(weights='distance'))
        ])
        clf = TransformedTargetClassifier(classifier=pipeline, transformer=get_target_transformer(y))
        param = {
                'classifier__model__p':[1,2]
        }
        grid = GridSearchCV(clf, param, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity and not get_model:
            print("Applying KNN...")
        grid.fit(x, y)
        if get_model :
            model = grid.best_estimator_
            return model, grid.best_score_, "KNN"
        return grid.best_score_
        
    def VotingClassifier_model(base_models, weights) :
        vc = VotingClassifier(estimators=base_models, weights=weights, voting="soft")
        grid = GridSearchCV(vc, param_grid={}, cv=cv, scoring=score, return_train_score=True, n_jobs=-1)
        if verbosity:
            print("Applying VotingClassifier...")
        grid.fit(x,y)
        return grid.best_estimator_, grid.best_score_, f"VotingClassifier{base_models}"
	
    if verbosity:
            print("Searching for the best fit...")
	
    func_map = {
        "LogisticRegression":lr_model,
        "DecisionTree":dt_model,
        "RandomForest":rf_model,
        "GaussianNB":GaussianNB_model,
        "BernoulliNB":BernoulliNB_model,
        "CategoricalNB":CategoricalNB_model,
        "CatBoost":CatBoost_model,
        "XGBoost":XGBoost_model,
        "Ripper":Ripper_model,
        "SVC":SVC_model,
        "KNN":KNN_model
    }

    if (whitelisted_algorithms=="auto") :
        if (x_type=="numeric") :
            whitelisted_algorithms = ["LogisticRegression", "DecisionTree", "RandomForest", "GaussianNB", "CatBoost", "XGBoost", "Ripper", "SVC", "KNN"]
        elif (x_type=="categorical") :
            whitelisted_algorithms = ["LogisticRegression", "DecisionTree", "RandomForest", "BernoulliNB", "CategoricalNB", "CatBoost", "XGBoost", "Ripper", "SVC", "KNN"]
        elif (x_type=="mixed") :
            whitelisted_algorithms = ["LogisticRegression", "DecisionTree", "RandomForest", "CatBoost", "XGBoost", "Ripper", "SVC", "KNN"]
        else :
            raise ValueError('x_type can only accept "numeric", "categorical", "mixed" values')

    dict_map = {}
    base_estimators = []
    weights = []
    list_of_tuples = []

    for element in whitelisted_algorithms :
        dict_map[element] = func_map[element]()

    dict_map = dict(sorted(dict_map.items(), key=lambda x : x[1], reverse=True))

    for e in dict_map :
        if (len(whitelisted_algorithms)>2 and dict_map[e] >= list(dict_map.values())[2]) :
            list_of_tuples.append(func_map[e](True))
            base_estimators.append((e, list_of_tuples[len(list_of_tuples)-1][0]))
            weights.append(dict_map[e])

    if (enable_votingClassifier and len(whitelisted_algorithms)>2) :
        normalized_weights = [w / sum(weights) for w in weights]
        m, v, n = VotingClassifier_model(base_estimators, normalized_weights)
    else :
        m, v, n = (0, 0, 0)
    
    if verbosity:
        print("Applying the best fit...")

    if (v < list(dict_map.values())[0]) :
        if len(list_of_tuples)>0:
            m, v, n = list_of_tuples[0]
        else:
            m, v, n = func_map[list(dict_map.keys())[0]](True)

    return _AutofitWrapper(m, v, n, x)



#=========================== tests ===========================#

from sklearn.datasets import load_iris
def test_trainer_kit() :
    """
    test trainer_kit with iris dataset
    """
    iris = load_iris()
    df = pd.DataFrame(iris.data, columns=iris.feature_names)
    df['target'] = iris.target
    print(len(df))
    df.tail(5)
    df['target'].unique()
    x = df.drop('target', axis=1)
    y = df['target']
    r = autofit_classification_model(x,y,"numeric",3,random_state=50)
    print(r.score)
    print(r.classifier)
    print(r.model.predict([[6.5,3,5.2,2]]))
#test_trainer_kit()