"""
Modified version of https://github.com/joaopfonseca/game-of-recourse/blob/main/game_viz/scorers.py
New Models added: Ridge Regression, SVM, MLP, Linear Regression
"""

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import RidgeClassifier
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.model_selection import KFold
from sklearn.svm import SVC
from sklearn.calibration import CalibratedClassifierCV
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import Pipeline
from tqdm import tqdm
from sklearn.metrics import accuracy_score, roc_auc_score, f1_score
from scipy.special import expit

class LinearReg:
    def __init__(self, ignore_features=None):
        """
        Parameters:
        - ignore_features: A list of features to ignore during training and prediction.
        """
        self.ignore_features = ignore_features if ignore_features is not None else []
        self.model = LinearRegression()

    def _get_X(self, X):
        # Remove the `ignore_features` from X
        return X.drop(columns=self.ignore_features, errors='ignore')

    def fit(self, X, y):
        # Filter out the `ignore_features` from X
        X_filtered = self._get_X(X)

        kf = KFold(n_splits=10)
        MSE_bag, R2_bag = [], []

        model = self.model
        model_name = "Linear Regression"
        print(model_name)
        for n, (train_index, test_index) in tqdm(enumerate(kf.split(X_filtered, y))):
            x_train, y_train = X_filtered.iloc[train_index], y.iloc[train_index]
            x_test, y_test = X_filtered.iloc[test_index], y.iloc[test_index]
            pipeline = Pipeline(steps=[('estimator', model)])
            pipeline.fit(x_train, y_train)
            y_pred = pipeline.predict(x_test)

            # For regression, we use MSE and R2 instead of classification metrics
            MSE_bag.append(mean_squared_error(y_test, y_pred))
            R2_bag.append(r2_score(y_test, y_pred))

        print(f'Mean Squared Error (MSE): {np.mean(MSE_bag):.4f} +/- {np.std(MSE_bag):.4f}')
        print(f'R2 Score: {np.mean(R2_bag):.4f} +/- {np.std(R2_bag):.4f}')

        # Store model coefficients
        self.coef_ = self.model.coef_
        self.intercept_ = self.model.intercept_
        return self

    def predict(self, X):
        """Predict the target value based on learned coefficients."""
        # Get continuous predictions (no thresholding needed for regression)
        return self.predict_proba(X)

    def predict_proba(self, X):
        """Return the predicted continuous values from the model."""
        X_filtered = self._get_X(X)
        return self.model.predict(X_filtered)

class LogReg:
    def __init__(self, threshold=0.5, ignore_features=None, intercept=0):
        """
        Parameters:
        - threshold: The threshold to use for prediction.
        - ignore_features: A list of features to ignore during training and prediction.
        """
        self.threshold = threshold
        # Set ignore_features to an empty list if None is provided
        self.ignore_features = ignore_features if ignore_features is not None else []
        self.intercept = intercept
        self.model = LogisticRegression(random_state=42)

    def _get_X(self, X):
        # Remove the 'groups' feature (or any other feature you want to ignore)
        return X.drop(columns=self.ignore_features, errors='ignore')

    def fit(self, X, y):
        # Filter out the `ignore_feature` from X
        X_filtered = self._get_X(X)

        kf = KFold(n_splits=10)
        AUC_bag, ACC_bag, F1_bag = [], [], []

        model = self.model
        model_name = "Logistic Regression"
        print(model_name)
        for n, (train_index, test_index) in tqdm(enumerate(kf.split(X_filtered, y))):
            x_train, y_train = X_filtered.iloc[train_index], y.iloc[train_index]
            x_test, y_test = X_filtered.iloc[test_index], y.iloc[test_index]
            pipeline = Pipeline(steps=[('estimator', model)])
            pipeline.fit(x_train, y_train["target"])
            y_pred = pipeline.predict(x_test)
            y_pred_prob = pipeline.predict_proba(x_test)
            AUC_bag.append(roc_auc_score(y_test["target"], y_pred_prob[:, 1]))
            ACC_bag.append(accuracy_score(y_test["target"], y_pred))
            F1_bag.append(f1_score(y_test["target"], y_pred))

        print(f'AUC: {np.mean(AUC_bag):.4f} +/- {np.std(AUC_bag):.4f}')
        print(f'ACC: {np.mean(ACC_bag):.4f} +/- {np.std(ACC_bag):.4f}')
        print(f'F1: {np.mean(F1_bag):.4f} +/- {np.std(F1_bag):.4f}')
        self.intercept_ = self.model.intercept_
        self.coef_ = self.model.coef_
        return self


    def predict(self, X):
        """Predict the target value based on learned probabilities."""
        # Get probabilities from `predict_proba`
        proba = self.predict_proba(X)
        # Use only the probability for class 1
        return (proba[:, 1] > self.threshold).astype(int).squeeze()


    def predict_proba(self, X):
        """Return the predicted probabilities from the model."""
        X_filtered = self._get_X(X)
        return self.model.predict_proba(X_filtered)
    
class RidgeReg:
    def __init__(self, threshold=0.5, ignore_features=None, intercept=0):
        """
        Parameters:
        - threshold: The threshold to use for prediction.
        - ignore_features: A list of features to ignore during training and prediction.
        """
        self.threshold = threshold
        self.ignore_features = ignore_features if ignore_features is not None else []
        self.intercept = intercept
        self.model = RidgeClassifier(alpha=1e-6, random_state=42)


    def _get_X(self, X):
        # Remove the 'groups' feature (or any other feature you want to ignore)
        return X.drop(columns=self.ignore_features, errors='ignore')

    def fit(self, X, y):
        # Filter out the `ignore_features` from X
        X_filtered = self._get_X(X)

        kf = KFold(n_splits=10)
        AUC_bag, ACC_bag, F1_bag = [], [], []

        model = self.model
        model_name = "Ridge Classifier"
        print(model_name)
        for n, (train_index, test_index) in tqdm(enumerate(kf.split(X_filtered, y))):
            x_train, y_train = X_filtered.iloc[train_index], y.iloc[train_index]
            x_test, y_test = X_filtered.iloc[test_index], y.iloc[test_index]
            pipeline = Pipeline(steps=[('scaler', StandardScaler()), ('estimator', model)])
            pipeline.fit(x_train, y_train["target"])
            y_pred = pipeline.predict(x_test)
            # Use sigmoid transformation on decision function for probabilities
            y_pred_prob = pipeline.decision_function(x_test)
            AUC_bag.append(roc_auc_score(y_test["target"], expit(y_pred_prob)))
            ACC_bag.append(accuracy_score(y_test["target"], y_pred))
            F1_bag.append(f1_score(y_test["target"], y_pred))

        print(f'AUC: {np.mean(AUC_bag):.4f} +/- {np.std(AUC_bag):.4f}')
        print(f'ACC: {np.mean(ACC_bag):.4f} +/- {np.std(ACC_bag):.4f}')
        print(f'F1: {np.mean(F1_bag):.4f} +/- {np.std(F1_bag):.4f}')
        self.intercept_ = pipeline.named_steps['estimator'].intercept_
        self.coef_ = pipeline.named_steps['estimator'].coef_
        return self

    def predict(self, X):
        """Predict the target value based on learned probabilities."""
        # Get probabilities from `predict_proba`
        proba = self.predict_proba(X)
        # Use the probability for the positive class (last column)
        return (proba[:, 1] > self.threshold).astype(int).squeeze()

    def predict_proba(self, X):
        """Return the predicted probabilities from the model."""
        X_filtered = self._get_X(X)
        decision_function = self.model.decision_function(X_filtered)
        positive_proba = expit(decision_function)  # Sigmoid transformation
        negative_proba = 1 - positive_proba
        # Stack negative and positive probabilities to form a 2D array
        return np.vstack((negative_proba, positive_proba)).T


class ModelSVM:
    def __init__(self, threshold=0.5, ignore_features=None, intercept=0, calibrate=False):
        """
        Parameters:
        - threshold: The threshold to use for prediction.
        - ignore_features: A list of features to ignore during training and prediction.
        - calibrate: Whether to calibrate the probabilities of the model.
        """
        self.threshold = threshold
        self.ignore_features = ignore_features if ignore_features is not None else []
        self.intercept = intercept
        self.calibrate = calibrate
        # Use linear kernel to make coefficients available
        self.model = SVC(kernel='linear', probability=True, random_state=42)

    def _get_X(self, X):
        """Filter out the ignored features."""
        return X.drop(columns=self.ignore_features, errors='ignore')

    def fit(self, X, y):
        """Train the model with K-fold cross-validation and optional probability calibration."""
        X_filtered = self._get_X(X)

        kf = KFold(n_splits=10)
        AUC_bag, ACC_bag, F1_bag = [], [], []

        model = self.model
        model_name = "SVM"
        print(model_name)

        for n, (train_index, test_index) in tqdm(enumerate(kf.split(X_filtered, y))):
            x_train, y_train = X_filtered.iloc[train_index], y.iloc[train_index]
            x_test, y_test = X_filtered.iloc[test_index], y.iloc[test_index]

            # Fit the model and optionally calibrate it
            pipeline = Pipeline(steps=[('estimator', model)])
            pipeline.fit(x_train, y_train["target"])

            # If calibration is enabled, calibrate the model
            if self.calibrate:
                pipeline = CalibratedClassifierCV(pipeline, method='sigmoid', cv='prefit')
                pipeline.fit(x_train, y_train["target"])

            y_pred = pipeline.predict(x_test)
            y_pred_prob = pipeline.predict_proba(x_test)

            AUC_bag.append(roc_auc_score(y_test["target"], y_pred_prob[:, 1]))
            ACC_bag.append(accuracy_score(y_test["target"], y_pred))
            F1_bag.append(f1_score(y_test["target"], y_pred))

        print(f'AUC: {np.mean(AUC_bag):.4f} +/- {np.std(AUC_bag):.4f}')
        print(f'ACC: {np.mean(ACC_bag):.4f} +/- {np.std(ACC_bag):.4f}')
        print(f'F1: {np.mean(F1_bag):.4f} +/- {np.std(F1_bag):.4f}')
        
        self.intercept_ = self.model.intercept_
        self.coef_ = self.model.coef_

        # If calibration was applied, use the calibrated model's parameters
        if self.calibrate:
            self.model = pipeline
        return self

    def predict(self, X):
        """Predict the target value based on learned probabilities."""
        proba = self.predict_proba(X)
        return (proba[:, 1] > self.threshold).astype(int).squeeze()

    def predict_proba(self, X):
        """Return the predicted probabilities from the model."""
        X_filtered = self._get_X(X)
        return self.model.predict_proba(X_filtered)
    
class ModelMLP:
    def __init__(self, threshold=0.5, ignore_features=None, intercept=0):
        """
        Parameters:
        - threshold: The threshold to use for prediction.
        - ignore_features: A list of features to ignore during training and prediction.
        """
        self.threshold = threshold
        # Set ignore_features to an empty list if None is provided
        self.ignore_features = ignore_features if ignore_features is not None else []
        self.intercept = intercept
        self.model = MLPClassifier(random_state=42)

    def _get_X(self, X):
        # Remove the 'groups' feature (or any other feature you want to ignore)
        return X.drop(columns=self.ignore_features, errors='ignore')

    def fit(self, X, y):
        # Filter out the `ignore_feature` from X
        X_filtered = self._get_X(X)

        kf = KFold(n_splits=10)
        AUC_bag, ACC_bag, F1_bag = [], [], []

        model = self.model
        model_name = "MLP"
        print(model_name)
        for n, (train_index, test_index) in tqdm(enumerate(kf.split(X_filtered, y))):
            x_train, y_train = X_filtered.iloc[train_index], y.iloc[train_index]
            x_test, y_test = X_filtered.iloc[test_index], y.iloc[test_index]
            pipeline = Pipeline(steps=[('estimator', model)])
            pipeline.fit(x_train, y_train["target"])
            y_pred = pipeline.predict(x_test)
            y_pred_prob = pipeline.predict_proba(x_test)
            AUC_bag.append(roc_auc_score(y_test["target"], y_pred_prob[:, 1]))
            ACC_bag.append(accuracy_score(y_test["target"], y_pred))
            F1_bag.append(f1_score(y_test["target"], y_pred))

        print(f'AUC: {np.mean(AUC_bag):.4f} +/- {np.std(AUC_bag):.4f}')
        print(f'ACC: {np.mean(ACC_bag):.4f} +/- {np.std(ACC_bag):.4f}')
        print(f'F1: {np.mean(F1_bag):.4f} +/- {np.std(F1_bag):.4f}')

        return self


    def predict(self, X):
        """Predict the target value based on learned probabilities."""
        # Get probabilities from `predict_proba`
        proba = self.predict_proba(X)
        # Use only the probability for class 1
        return (proba[:, 1] > self.threshold).astype(int).squeeze()


    def predict_proba(self, X):
        """Return the predicted probabilities from the model."""
        X_filtered = self._get_X(X)
        return self.model.predict_proba(X_filtered)