import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures


class LinearRegressionModel(BaseEstimator, TransformerMixin):
    def __init__(
        self,
        target_variable: str = '',
        variable_columns: list | None = None,
        model_params: dict | None = None,
        clipping_max: float | None = None,
        clipping_min: float | None = None,
        weights: dict | None = None,
        degree: int = 1,
        interaction_only: bool = False,
        verbose: bool = False,
    ):
        """
        Linear Regression Model for Time Series Analysis.

        Args:
            target_variable (str): The target variable name
            variable_columns (list): The input columns names in a list
            model_params (dict): The parameters used for training the model \\
            clipping (dict): The lower and upper limits for the target variable to be clipped \\
                *Format: {'min': min_value, 'max': max_value}*
            weights (dict): The weights for the Linear Regression model \\
                *Format: {'variable_name': weight}*
            degree (int): The degree of the polynomial features
            interaction_only (bool): Whether to include interaction terms only
            verbose (bool): Whether to print verbose output

        Returns:
            LinearRegressionModel: The prediction model object
        """
        self.target_variable: str = target_variable
        self.variable_columns: list = variable_columns if variable_columns else []
        self.model_params: dict = model_params if model_params else {}
        self.regr = LinearRegression()
        self.q1_target: float | None = None
        self.q3_target: float | None = None
        self.weights: dict = weights if weights else {}
        self.degree: int = degree
        self.interaction_only: bool = interaction_only
        self.poly: PolynomialFeatures | None = None
        self.verbose: bool = verbose

        self.clipping: dict = {}

        if clipping_max is not None:
            self.clipping['max'] = clipping_max
        if clipping_min is not None:
            self.clipping['min'] = clipping_min

    def get_regressor(self) -> LinearRegression:
        """
        Get LinearRegression object

        Returns:
            LinearRegression objetc
        """
        return self.regr

    def create_poly_features(self, input_data: pd.DataFrame, fit: bool = False) -> pd.DataFrame:
        """
        Function to create polynomial features

        Args:
            input_data (pandas.DataFrame): The data used to create the polynomial features
            fit (bool): Whether to fit the polynomial features creator or just transform

        Returns:
            pandas.DataFrame: The data with the polynomial features
        """
        original_index = input_data.index
        original_shape = input_data.shape

        # Create polynomial features
        interaction_only = self.interaction_only

        degree = self.degree
        if fit:
            if self.verbose:
                print(f'Fitting polynomial features with degree {degree}')
            self.poly = PolynomialFeatures(
                degree=degree, interaction_only=interaction_only, include_bias=False
            )
            input_data = self.poly.fit_transform(input_data)
        else:
            if self.verbose:
                print(f'Transforming polynomial features with degree {degree}')
            if self.poly is not None:
                input_data = self.poly.transform(input_data)
            else:
                raise ValueError(
                    'Polynomial transformer is not fitted. Call fit() before predict() when degree > 1.'
                )

        # Create new columns names
        modified_columns = self.poly.get_feature_names_out(self.variable_columns)

        input_data = pd.DataFrame(input_data, columns=modified_columns, index=original_index)

        if self.verbose:
            print(f'Data shape before transformation:   {original_shape}')
            print(f'Data shape after transformation:    {input_data.shape}')

        return input_data

    def fit(self, input_data: pd.DataFrame) -> 'LinearRegressionModel':
        """
        Function to fit the model

        Args:
            input_data (pandas.DataFrame): The data used to fit the Linear Regression model

        Returns:
            LinearRegressionModel: The prediction model object
        """
        if self.verbose:
            # Display the header
            text = 'INITIATING LINEAR REGRESSION MODEL FIT'
            print('\n' + '-' * len(text))
            print(text)
            print('-' * len(text) + '\n')

        # Basic validations and inference
        if not self.target_variable:
            raise ValueError("'target_variable' must be set before calling fit().")

        if not self.variable_columns:
            # Infer all columns except target
            self.variable_columns = [c for c in input_data.columns if c != self.target_variable]

        # Validate required columns exist
        missing_features = [c for c in self.variable_columns if c not in input_data.columns]
        if missing_features:
            raise ValueError(
                'Training data is missing required feature columns: ' + ', '.join(missing_features)
            )

        if self.target_variable not in input_data.columns:
            # Let pandas raise KeyError with the column name to satisfy tests too
            raise KeyError(repr(self.target_variable))

        X_train = input_data[self.variable_columns]
        y_train = input_data[self.target_variable]

        if self.verbose:
            print(f'X_train shape: {X_train.shape}')
            print(f'y_train shape: {y_train.shape}')

        # Create polynomial features
        if self.degree > 1:
            X_train = self.create_poly_features(X_train, fit=True)

        # Handle infinite values
        X_train = X_train.replace([np.inf, -np.inf], np.nan)

        # Drop columns with all null values
        all_null_cols = X_train.columns[X_train.isnull().all()].tolist()
        X_train = X_train.drop(columns=all_null_cols)

        if self.verbose and all_null_cols:
            print(f'{len(all_null_cols)} columns with all null values were dropped.')

        # Get mask of rows with any NaN values
        mask = X_train.notna().all(axis=1)
        X_train = X_train[mask]
        y_train = y_train[mask]
        if self.verbose:
            print(f'X_train shape after removing NaN rows: {X_train.shape}')
            print(f'y_train shape after removing NaN rows: {y_train.shape}')

        self.q1_target = y_train.quantile(0.25)
        self.q3_target = y_train.quantile(0.75)

        # Fit the model
        self.regr.fit(X_train, y_train)
        if self.verbose:
            print('Model training completed successfully.')

        # Get the weights using actual trained feature names (handles polynomial features)
        # ravel() so coef_ is 1D when y was a column (e.g. DataFrame), avoiding scalar conversion errors
        round_coef = np.round(self.regr.coef_, 3).ravel()
        round_intercept = np.round(self.regr.intercept_, 3).ravel()
        feature_names = list(getattr(self.regr, 'feature_names_in_', X_train.columns))

        # Save the weights (strict=True to ensure feature_names and coefficients align)
        weights = dict(zip(feature_names, [float(c) for c in round_coef], strict=True))
        weights = dict(sorted(weights.items(), key=lambda item: abs(item[1]), reverse=True))
        weights = {'Bias': float(round_intercept.flat[0]), **weights}
        self.weights = weights

        if self.verbose:
            print('\nModel Weights:')
            for key, value in weights.items():
                print(f'{key}: {value}')

        return self

    def predict(self, input_data: pd.DataFrame) -> np.ndarray:
        """
        Function to predict the target variable.
        If clipping is True, the predictions are clipped based on the target variable quartiles.

        Args:
            input_data (pandas.DataFrame): The data used to predict the target variable

        Returns:
            numpy.ndarray: The predicted target variable
        """
        if self.verbose:
            # Display the header
            text = 'INITIATING LINEAR REGRESSION MODEL PREDICTION'
            print('\n' + '-' * len(text))
            print(text)
            print('-' * len(text) + '\n')

        # Validate required columns exist for prediction
        missing_features = [c for c in self.variable_columns if c not in input_data.columns]
        if missing_features:
            raise ValueError(
                'Prediction data is missing required feature columns: '
                + ', '.join(missing_features)
            )

        X_test = input_data[self.variable_columns]
        if self.verbose:
            print(f'X_test shape: {X_test.shape}')

        # Create polynomial features
        if self.degree > 1:
            X_test = self.create_poly_features(X_test, fit=False)
        X_test = X_test.replace([np.inf, -np.inf], np.nan)

        model_features = self.regr.feature_names_in_
        # Ensure required model features exist
        missing_model_features = [c for c in model_features if c not in X_test.columns]
        if missing_model_features:
            raise ValueError(
                'Prediction data is missing required transformed feature columns: '
                + ', '.join(missing_model_features)
            )

        # Get mask of rows with any NaN values
        mask = X_test.notna().all(axis=1)
        X_test = X_test[mask]
        if self.verbose:
            print(f'X_test shape after removing NaN rows: {X_test.shape}')

        # For simplicity and to keep behavior consistent with prior tests, drop NaN rows before prediction
        y_pred = self.regr.predict(X_test[model_features])

        if self.clipping:
            for i in range(len(y_pred)):
                if y_pred[i] > self.clipping['max']:
                    y_pred[i] = self.q3_target
                elif y_pred[i] < self.clipping['min']:
                    y_pred[i] = self.q1_target

        return y_pred