import numpy as np

class LinearRegressionOLS:
    def __init__(self):
        self.m = None
        self.b = None

    def fit(self, X_train, y_train):
        try:
            if X_train.ndim != 1 or y_train.ndim != 1:
                raise ValueError("Only 1D arrays supported for this simple linear regression")

            num = 0
            den = 0

            for i in range(X_train.shape[0]):
                num += (X_train[i] - X_train.mean()) * (y_train[i] - y_train.mean())
                den += (X_train[i] - X_train.mean())**2

            self.m = num / den
            self.b = y_train.mean() - (self.m * X_train.mean())

        except Exception as e:
            print(f"[ERROR in fit]: {e}")
            self.m = None
            self.b = None
    
    def predict(self, X_test):
        return (self.m * X_test) + self.b

    def checkScore(self, X_test, y_test):

        try:
            y_pred = self.predict(X_test)
            ss_res = sum((y_test - y_pred) ** 2)
            ss_tot = sum((y_test - y_test.mean()) ** 2)
            return 1 - (ss_res / ss_tot)
        except Exception as e:
            print(f"[ERROR in checkScore]: {e}")
            return None
        
    def plotGraph(self, X_test, y_test):
    
        import matplotlib.pyplot as plt
        plt.scatter(X_test, y_test, color = "blue")
        y_pred = self.predict(X_test)
        plt.plot(X_test, y_pred, color = "red")
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.legend(["Data Points", "Regression Line"])
        plt.show()
    
    def getEquation(self):
        return f'y_hat = {self.m} * X + {self.b}'

class MLRegressionOLS:
    def __init__(self):
        self.coef_ = None
        self.intercept_ = None

    def fit(self, X_train, y_train):
        X_train = np.insert(X_train, 0, 1, axis=1)
        betas = np.linalg.inv(np.dot(X_train.T, X_train)).dot(X_train.T).dot(y_train)
        self.intercept_ = betas[0]
        self.coef_ = betas[1:]

    def predict(self, X_test):
        y_pred = np.dot(X_test, self.coef_) + self.intercept_
        return y_pred

    def checkScore(self, X_test, y_test):
        try:
            y_pred = self.predict(X_test)
            ss_res = np.sum((y_test - y_pred) ** 2)
            ss_tot = np.sum((y_test - np.mean(y_test)) ** 2)
            return 1 - (ss_res / ss_tot)
        except Exception as e:
            print(f"[ERROR in checkScore]: {e}")
            return None

    def plotGraph(self, X_test, y_test):
        import plotly.graph_objs as go
        
        # Generate meshgrid from X_test
        x_range = np.linspace(X_test[:, 0].min(), X_test[:, 0].max(), 30)
        y_range = np.linspace(X_test[:, 1].min(), X_test[:, 1].max(), 30)
        X1_grid, X2_grid = np.meshgrid(x_range, y_range)
        
        # Flatten mesh and predict Z (Y values)
        mesh_input = np.c_[X1_grid.ravel(), X2_grid.ravel()]
        Z = self.predict(mesh_input).reshape(X1_grid.shape)

        # Plot original points
        fig = go.Figure()

        fig.add_trace(go.Scatter3d(
            x=X_test[:, 0], y=X_test[:, 1], z=y_test,
            mode='markers', marker=dict(color='blue', size=5),
            name='Data Points'
        ))

        # Plot regression plane
        fig.add_trace(go.Surface(
            x=X1_grid, y=X2_grid, z=Z,
            colorscale='Viridis',
            name='Regression Plane',
            opacity=0.7
        ))

        fig.update_layout(scene=dict(
            xaxis_title='X1',
            yaxis_title='X2',
            zaxis_title='Y'
        ))

        fig.show()

    def getEquation(self):
        return f'y_hat = {self.intercept_} + {self.coef_} * X'
    
    
class LinearRegression_With_BGD:
    def __init__(self, epochs, learning_rate):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.m = 0
        self.b = 0


    def fit(self, X_train, y_train):
        n = X_train.shape[0]
        
        for _ in range(self.epochs):
            y_pred = self.m * X_train + self.b
            error = y_train - y_pred
            
            der_m = (-2 / n) * (X_train * error).sum()
            der_b = (-2 / n) * error.sum()

            self.m -= self.learning_rate * der_m
            self.b -= self.learning_rate * der_b

    def predict(self, X_test):
        return (self.m * X_test) + self.b

    def checkScore(self, X_test, y_test):
        try:
            y_pred = self.predict(X_test)
            ss_res = np.sum((y_test - y_pred) ** 2)
            ss_tot = np.sum((y_test - np.mean(y_test)) ** 2)
            return 1 - (ss_res / ss_tot)
        except Exception as e:
            print(f"[ERROR in checkScore]: {e}")
            return None
    
    def getEquation(self):
        print(f'y_hat = {self.m} * X + {self.b}')

    def plotGraph(self, X_test, y_test):
        import matplotlib.pyplot as plt
        plt.scatter(X_test, y_test, color='blue')
        y_pred = self.predict(X_test)
        plt.plot(X_test, y_pred, color='red')
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.legend(["Data Points", "Regression Line"])
        plt.show()
    

class MLRegressionBGD:
    def __init__(self, epochs=100, learning_rate=0.01):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.coef_ = None
        self.intercept_ = None

    def fit(self, X_train, y_train):
        n = X_train.shape[0]
        X_train = np.insert(X_train, 0, 1, axis=1)
        self.coef_ = np.zeros(X_train.shape[1])
        for _ in range(self.epochs):
            y_pred = X_train.dot(self.coef_)
            error = y_train - y_pred
            gradient = - (2/n) * X_train.T.dot(error)
            self.coef_ -= self.learning_rate * gradient

    def predict(self, X_test):
        X_test = np.insert(X_test, 0, 1, axis=1)
        return X_test.dot(self.coef_)

    def checkScore(self, X_test, y_test):
        try:
            y_pred = self.predict(X_test)
            ss_res = np.sum((y_test - y_pred) ** 2)
            ss_tot = np.sum((y_test - np.mean(y_test)) ** 2)
            return 1 - (ss_res / ss_tot)
        except Exception as e:
            print(f"[ERROR in checkScore]: {e}")
            return None
    
    def plotGraph(self, X_test, y_test):
        import plotly.graph_objs as go
        
        # Generate meshgrid from X_test
        x_range = np.linspace(X_test[:, 0].min(), X_test[:, 0].max(), 30)
        y_range = np.linspace(X_test[:, 1].min(), X_test[:, 1].max(), 30)
        X1_grid, X2_grid = np.meshgrid(x_range, y_range)
        
        # Flatten mesh and predict Z (Y values)
        mesh_input = np.c_[X1_grid.ravel(), X2_grid.ravel()]
        Z = self.predict(mesh_input).reshape(X1_grid.shape)

        # Plot original points
        fig = go.Figure()

        fig.add_trace(go.Scatter3d(
            x=X_test[:, 0], y=X_test[:, 1], z=y_test,
            mode='markers', marker=dict(color='blue', size=5),
            name='Data Points'
        ))

        # Plot regression plane
        fig.add_trace(go.Surface(
            x=X1_grid, y=X2_grid, z=Z,
            colorscale='Viridis',
            name='Regression Plane',
            opacity=0.7
        ))

        fig.update_layout(scene=dict(
            xaxis_title='X1',
            yaxis_title='X2',
            zaxis_title='Y'
        ))

        fig.show()
    
    def getEquation(self):
        return f'y_hat = {self.coef_[0]} + {self.coef_[1:]} * X'

class LinearRegressionSGD:
    def __init__(self, epochs, learning_rate):
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.m = 0
        self.b = 0

    def fit(self, X_train, y_train):
        for _ in range(self.epochs):
            for _ in range(len(X_train)):
                idx = np.random.randint(0, len(X_train))
                error = y_train[idx] - (self.m * X_train[idx] + self.b)
                self.m += self.learning_rate * 2 * error * X_train[idx]
                self.b += self.learning_rate * 2 * error

    
    def predict(self, X_test):
        return (self.m * X_test) + self.b

    def checkScore(self, X_test, y_test):
        try:
            y_pred = self.predict(X_test)
            ss_res = np.sum((y_test - y_pred) ** 2)
            ss_tot = np.sum((y_test - np.mean(y_test)) ** 2)
            return 1 - (ss_res / ss_tot)
        except Exception as e:
            print(f"[ERROR in checkScore]: {e}")
            return None
    
    def getEquation(self):
        print(f'y_hat = {self.m} * X + {self.b}')

    def plotGraph(self, X_test, y_test):
        import matplotlib.pyplot as plt
        plt.scatter(X_test, y_test, color='blue')
        y_pred = self.predict(X_test)
        plt.plot(X_test, y_pred, color='red')
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.legend(["Data Points", "Regression Line"])
        plt.show()

    
class MLRegressionSGD:
    
    def __init__(self,learning_rate=0.01,epochs=100):
        
        self.coef_ = None
        self.intercept_ = None
        self.lr = learning_rate
        self.epochs = epochs
        
    def fit(self, X_train, y_train):
        # Add bias column (1s) at front
        X_train = np.insert(X_train, 0, 1, axis=1)
        self.coef_ = np.ones(X_train.shape[1])  # include bias in coef_

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

    def predict(self, X_test):
        X_test = np.insert(X_test, 0, 1, axis=1)
        return np.dot(X_test, self.coef_)

    def checkScore(self, X_test, y_test):
        y_pred = self.predict(X_test)
        ss_total = np.sum((y_test - np.mean(y_test))**2)
        ss_res = np.sum((y_test - y_pred)**2)
        return 1 - (ss_res / ss_total)

    def plotGraph(self, X_test, y_test):
        import plotly.graph_objs as go
        
        # Generate meshgrid from X_test
        x_range = np.linspace(X_test[:, 0].min(), X_test[:, 0].max(), 30)
        y_range = np.linspace(X_test[:, 1].min(), X_test[:, 1].max(), 30)
        X1_grid, X2_grid = np.meshgrid(x_range, y_range)
        
        # Flatten mesh and predict Z (Y values)
        mesh_input = np.c_[X1_grid.ravel(), X2_grid.ravel()]
        Z = self.predict(mesh_input).reshape(X1_grid.shape)

        # Plot original points
        fig = go.Figure()

        fig.add_trace(go.Scatter3d(
            x=X_test[:, 0], y=X_test[:, 1], z=y_test,
            mode='markers', marker=dict(color='blue', size=5),
            name='Data Points'
        ))

        # Plot regression plane
        fig.add_trace(go.Surface(
            x=X1_grid, y=X2_grid, z=Z,
            colorscale='Viridis',
            name='Regression Plane',
            opacity=0.7
        ))

        fig.update_layout(scene=dict(
            xaxis_title='X1',
            yaxis_title='X2',
            zaxis_title='Y'
        ))

        fig.show()

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

   
    
    

        
    

    
    





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

    def split_data(self, ratio = 0.8, randomState = 42):
        try:
            # Always use self.target_column instead of argument
            if self.target_column is None or self.target_column not in self.dataframe.columns:
                raise ValueError("Please provide a valid target_column present in the dataframe.")
            train = self.dataframe.sample(frac=ratio, random_state=randomState)
            test = self.dataframe.drop(train.index)

            self.X_train = train.drop(columns=[self.target_column])
            self.y_train = train[self.target_column]
            self.X_test = test.drop(columns=[self.target_column])
            self.y_test = test[self.target_column]

        except Exception as e:
            print(f"Error in split_data: {e}")
            self.X_train, self.X_test, self.y_train, self.y_test = None, None, None, None
    
    def Linear_Regression_OLS(self, get_equation = False, plot = False, accuracy = True):
        lr_ols = LinearRegressionOLS()
        # Select the first column for simple linear regression
        X_train_1d = self.X_train.iloc[:, 0].values
        X_test_1d = self.X_test.iloc[:, 0].values
        y_train_1d = self.y_train.values
        y_test_1d = self.y_test.values
        lr_ols.fit(X_train_1d, y_train_1d)
        if (get_equation):
            print(lr_ols.getEquation())

        if (plot):
            lr_ols.plotGraph(X_test_1d, y_test_1d)

        if (accuracy):
            print(lr_ols.checkScore(X_test_1d, y_test_1d))
        
        elif (plot == False and accuracy == False):
            print("Model Trained")
    
    def MLinear_Regression_OLS(self, get_equation=False, plot=False, accuracy=True):
        mlr_ols = MLRegressionOLS()
        X_train_nd = self.X_train.values
        X_test_nd = self.X_test.values
        y_train_nd = self.y_train.values
        y_test_nd = self.y_test.values
        mlr_ols.fit(X_train_nd, y_train_nd)
        if get_equation:
            print(mlr_ols.getEquation())

        if plot:
            mlr_ols.plotGraph(X_test_nd, y_test_nd)

        if accuracy:
            print(mlr_ols.checkScore(X_test_nd, y_test_nd))

        elif not plot and not accuracy:
            print("Model Trained")

    def MLinear_Regression_SGD(self, epochs = 100, learning_rate = 0.01, plot = False, accuracy = True, get_equation = False):
        mlr_sgd = MLRegressionSGD(epochs=epochs, learning_rate=learning_rate)
        
        X_train_nd = self.X_train.values
        X_test_nd = self.X_test.values
        y_train_nd = self.y_train.values
        y_test_nd = self.y_test.values

        mlr_sgd.fit(X_train_nd, y_train_nd)

        if get_equation:
            print(mlr_sgd.getEquation())

        if plot:
            mlr_sgd.plotGraph(X_test_nd, y_test_nd)

        if accuracy:
            print(mlr_sgd.checkScore(X_test_nd, y_test_nd))

        elif not plot and not accuracy:
            print("Model Trained")

    def Linear_Regression_BGD(self, epochs = 100, learning_rate = 0.01, plot = False, accuracy = True, get_equation = False):
        lr_bgd = LinearRegression_With_BGD(epochs=epochs, learning_rate=learning_rate)
        X_train_1d = self.X_train.iloc[:, 0].values
        X_test_1d = self.X_test.iloc[:, 0].values
        y_train_1d = self.y_train.values
        y_test_1d = self.y_test.values
        lr_bgd.fit(X_train_1d, y_train_1d)
        if get_equation:
            print(lr_bgd.getEquation())

        if plot:
            lr_bgd.plotGraph(X_test_1d, y_test_1d)

        if accuracy:
            print(lr_bgd.checkScore(X_test_1d, y_test_1d))

        elif not plot and not accuracy:
            print("Model Trained")

    def MLinear_Regression_BGD(self, epochs=100, learning_rate=0.01, plot=False, accuracy=True, get_equation=False):
        mlr_bgd = MLRegressionBGD(epochs=epochs, learning_rate=learning_rate)
        X_train_nd = self.X_train.values
        X_test_nd = self.X_test.values
        y_train_nd = self.y_train.values
        y_test_nd = self.y_test.values
        mlr_bgd.fit(X_train_nd, y_train_nd)
        if get_equation:
            print(mlr_bgd.getEquation())

        if plot:
            mlr_bgd.plotGraph(X_test_nd, y_test_nd)

        if accuracy:
            print(mlr_bgd.checkScore(X_test_nd, y_test_nd))

        elif not plot and not accuracy:
            print("Model Trained")
        
    def Standard_Scale(self, features=None):
        try:
            if features is None:
                features = self.dataframe.columns.tolist()

                def standardized(X):
                    mean = X.mean(axis=0)
                    std = X.std(axis=0)
                    return (X - mean) / std

                self.dataframe[features] = standardized(self.dataframe[features])
                self.split_data()
          
        except Exception as e:
            print(f"Error in Standard_Scale: {e}")
    
    def Linear_Regression_SGD(self, epochs = 100, learning_rate = 0.01, plot=False, accuracy=True, get_equation=False):
        lr_sgd = LinearRegressionSGD(epochs=epochs, learning_rate=learning_rate)
        X_train_1d = self.X_train.iloc[:, 0].values
        X_test_1d = self.X_test.iloc[:, 0].values
        y_train_1d = self.y_train.values
        y_test_1d = self.y_test.values
        lr_sgd.fit(X_train_1d, y_train_1d)
        if get_equation:
            print(lr_sgd.getEquation())

        if plot:
            lr_sgd.plotGraph(X_test_1d, y_test_1d)

        if accuracy:
            print(lr_sgd.checkScore(X_test_1d, y_test_1d))

        elif not plot and not accuracy:
            print("Model Trained")

    def Extract_Data(self):
        return self.X_train, self.X_test, self.y_train, self.y_test
    
    def Set_DF(self, newDF, target_column):
        self.dataframe = newDF.copy()
        self.target_column = target_column
        self.split_data()

