import pandas as pd
import numpy as np
import pickle

class SmartTransformer:    
    @staticmethod
    def transform(X, y):
        X = X.values
        y = y.values

        if X.ndim == 1:
            X = X.reshape(-1, 1)
        return X, y

class Linear_Regression_OLS:
    def __init__(self):
        self.coef_ = None

    def fit(self, X_train, y_train):
        if X_train.ndim == 1:
            X_train = X_train.reshape(-1, 1)

        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.dot(np.linalg.inv(np.dot(X_train.T, X_train)), np.dot(X_train.T, y_train))

    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()

class Linear_regression_BGD:
    def __init__(self, epochs, learning_rate):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None

    def fit(self, X_train, y_train):
        if (X_train.ndim == 1):
            X_train = X_train.reshape(-1, 1)

        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.zeros( X_train.shape[1])
        n = X_train.shape[0]

        for _ in range(self.epochs):
            error = y_train - np.dot(X_train, self.coef_)
            slope = (-2/n) * (np.dot(X_train.T, error))
            self.coef_ -= (self.learning_rate * slope) 

    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()
    
    

class Linear_Regression_BGD_L1:
    def __init__(self, epochs, learning_rate, theLamda):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None
        self.thisLambda = theLamda

    def fit(self, X_train, y_train):
        if X_train.ndim == 1:
            X_train = X_train.reshape(-1, 1)
       
        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.zeros(X_train.shape[1])
        n = X_train.shape[0]

        for _ in range(self.epochs):
            error = y_train - np.dot(X_train, self.coef_)
            slope = (-2 / n) * np.dot(X_train.T, error) + self.thisLambda * np.sign(self.coef_)
            self.coef_ -= self.learning_rate * slope
    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()
    
class Linear_Regression_BGD_L2:
    def __init__(self, epochs, learning_rate, theLamda):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None
        self.thisLambda = theLamda

    def fit(self, X_train, y_train):
        if X_train.ndim == 1:
            X_train = X_train.reshape(-1, 1)
       
        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.zeros(X_train.shape[1])
        n = X_train.shape[0]
        for _ in range(self.epochs):
            error = y_train - (np.dot(X_train, self.coef_))
            slope = (-2/n) * (np.dot(X_train.T, error)) + (2 * self.thisLambda * (self.coef_))
            self.coef_ -= (self.learning_rate * slope)

    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()

class Linear_regression_SGD:
    def __init__(self, epochs, learning_rate):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None
    
    def fit(self, X_train, y_train):
        if X_train.ndim == 1:
            X_train = X_train.reshape(-1, 1)
       
        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.zeros(X_train.shape[1])

        for _ in range(self.epochs):
            for _ in range(X_train.shape[0]):
                idx = np.random.randint(0, X_train.shape[0])
                y_pred = np.dot(X_train[idx], self.coef_)
                error = y_train[idx] - y_pred
                slope = (-2) * (np.dot(X_train[idx].T, error))
                self.coef_ -= (self.learning_rate * slope) 

    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()

class Linear_regression_SGD_L1:
    def __init__(self, epochs, learning_rate, theLamda):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None
        self.thisLambda = theLamda
    
    def fit(self, X_train, y_train):
        if X_train.ndim == 1:
            X_train = X_train.reshape(-1, 1)
       
        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.zeros(X_train.shape[1])

        for _ in range(self.epochs):
            for _ in range(X_train.shape[0]):
                idx = np.random.randint(0, X_train.shape[0])
                y_pred = np.dot(X_train[idx], self.coef_)
                error = y_train[idx] - y_pred
                slope = (-2) * (np.dot(X_train[idx].T, error)) + (self.thisLambda * np.sign(self.coef_))
                self.coef_ -= (self.learning_rate * slope) 

    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()

class Linear_regression_SGD_L2:
    def __init__(self, epochs, learning_rate, theLamda):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None
        self.thisLambda = theLamda
    
    def fit(self, X_train, y_train):
        if X_train.ndim == 1:
            X_train = X_train.reshape(-1, 1)
       
        X_train = np.hstack((np.ones((X_train.shape[0], 1)), X_train))
        self.coef_ = np.zeros(X_train.shape[1])

        for _ in range(self.epochs):
            for _ in range(X_train.shape[0]):
                idx = np.random.randint(0, X_train.shape[0])
                y_pred = np.dot(X_train[idx], self.coef_)
                error = y_train[idx] - y_pred
                slope = (-2) * (np.dot(X_train[idx].T, error)) + (self.thisLambda * 2 * self.coef_)
                self.coef_ -= (self.learning_rate * slope) 

    def predict(self, X_test):
        pre = Predict_Return_Equation(X_test=X_test, coef_=self.coef_)
        return pre.predict()

    def returnScore(self, X_test, y_test):
        pre = Predict_Return_Equation(X_test=X_test, y_test=y_test, coef_=self.coef_)
        return pre.returnScore()

    def getEquation(self):
        pre = Predict_Return_Equation(coef_=self.coef_)
        return pre.getEquation()


    
class Predict_Return_Equation:
    def __init__(self, X_test=None, y_test=None, coef_=None):
        self.coef_ = coef_
        self.X_test = X_test
        self.y_test = y_test

    def predict(self):
        self.X_test = np.hstack((np.ones((self.X_test.shape[0], 1)), self.X_test))
        return np.dot(self.X_test, self.coef_)

    def returnScore(self):
        y_pred = self.predict()
        ss_res = np.sum((self.y_test - y_pred) ** 2)
        ss_tot = np.sum((self.y_test - np.mean(self.y_test)) ** 2)
        return 1 - (ss_res / ss_tot)

    def getEquation(self):
        intercept = round(self.coef_[0], 5)
        coefs = np.round(self.coef_[1:], 5)
        return f"y_hat = {intercept} + {coefs} * X"

        
class Lazy_Work:
    def __init__(self, df):
        self.dataframe = df.copy()
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
        self.target_column = None

    def fit_data(self, random_state=42, ratio=0.8, training_features=None, target_features=None, drop_features=None):
        # Ensure target_features is a list
        if target_features is None:
            raise ValueError("target_features cannot be None")
        elif isinstance(target_features, str):
            target_features = [target_features]
        elif not isinstance(target_features, list):
            target_features = list(target_features)
        
        # Validate only one target feature
        if len(target_features) > 1:
            raise ValueError("Only one target feature is allowed")
        
        self.target_column = target_features[0]
        
        # Set default training_features (all columns except target)
        if training_features is None:
            training_features = self.dataframe.columns.drop(self.target_column).tolist()
        # Convert training_features to list if not None
        elif isinstance(training_features, str):
            training_features = [training_features]
        else:
            training_features = list(training_features)
        
        # Remove features specified in drop_features
        if drop_features is not None:
            if isinstance(drop_features, str):
                drop_features = [drop_features]
            else:
                drop_features = list(drop_features)
            training_features = [f for f in training_features if f not in drop_features]
        
        # Filter columns (ensure target is included)
        df_filtered = self.dataframe[training_features + [self.target_column]]
        
        # Train-test split
        train = df_filtered.sample(frac=ratio, random_state=random_state)
        test = df_filtered.drop(train.index)
        
        # Assign features and target
        self.X_train = train[training_features]
        self.y_train = train[self.target_column]
        self.X_test = test[training_features]
        self.y_test = test[self.target_column]
        
        # Apply transformations
        self.X_train, self.y_train = SmartTransformer.transform(self.X_train, self.y_train)
        self.X_test, self.y_test = SmartTransformer.transform(self.X_test, self.y_test)


    def Standard_Scale(self, features=None):
        def standardized(X):
            mean = X.mean(axis=0)
            std = X.std(axis=0)
            return (X - mean) / std

        if features is None:
            features = self.dataframe.columns.tolist()

        self.dataframe[features] = standardized(self.dataframe[features])


    def save_model(self, model_obj, filename="model.pkl"):
        with open(filename, "wb") as f:
            pickle.dump(model_obj, f)

    def load_model(self, filename="model.pkl"):
        with open(filename, "rb") as f:
            return pickle.load(f)
    
    def doML(self, get_equation = False, model = None, method = None, penalty = None, epochs = 100, learning_rate = 0.01, lamda_ = 0.01):
        if (model is None or method is None):
            raise ValueError("Parameter can't be None")
        
        if (model == "lr"):
            if (penalty is None):
                if (method == "ols"):
                    lr = Linear_Regression_OLS()
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)
                elif (method == "bgd"):
                    lr = Linear_regression_BGD(epochs=epochs, learning_rate=learning_rate)
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)
                elif (method == "sgd"):
                    lr = Linear_regression_SGD(epochs=epochs, learning_rate=learning_rate)
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)
                
            
            elif (penalty == 'l1'):
                if (method == "bgd"):
                    lr = Linear_Regression_BGD_L1(epochs=epochs, learning_rate=learning_rate, theLamda=lamda_)
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)
                elif(method == "sgd"):
                    lr = Linear_regression_SGD_L1(epochs=epochs, learning_rate=learning_rate, theLamda=lamda_)
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)

            elif (penalty == 'l2'):
                if (method == "bgd"):
                    lr = Linear_Regression_BGD_L2(epochs=epochs, learning_rate=learning_rate, theLamda=lamda_)
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)
                elif (method == 'sgd'):
                    lr = Linear_regression_SGD_L2(epochs=epochs, learning_rate=learning_rate, theLamda=lamda_)
                    lr.fit(self.X_train, self.y_train)
                    if (get_equation):
                        print(lr.getEquation())
                    return lr.returnScore(self.X_test, self.y_test)

            
                
            
