from sklearn.metrics import r2_score
import numpy as np
from scipy.optimize import curve_fit
from scipy.stats import t


class NonLinearRegression:
    def __init__(self, initial_guess, T_ref):
        """
        Initialize the NonLinearRegression class with initial parameters.

        Parameters:
        - initial_guess: List of initial guesses for A' and Ea.
        - T_ref: Reference temperature
        """
        self.initial_guess = initial_guess
        self.T_ref = T_ref
        self.params_opt = None
        self.param_se = None
        self.conf_intervals = None
        self.r_squared = None

    @staticmethod
    def kinetic_rate_law(t, A_prime, Ea, T, T_ref):
        """
        Non-linear kinetic rate law function.

        Parameters:
        - t: Time (independent variable).
        - A_prime: Pre-exponential factor (parameter).
        - Ea: Activation energy (parameter).
        - T: Temperature data (array).
        - T_ref: Reference temperature.

        Returns:
        - Computed dependent variable.
        """
        R = 8.314  # Universal gas constant (J/(mol·K))
        return np.exp(-A_prime * np.exp(-Ea / R * (1 / T - 1 / T_ref)) * t)

    def fit(self, time_data, temp_data, c_c0_data):
        """
        Perform non-linear regression using global curve fitting.

        Parameters:
        - time_data: Time data (independent variable).
        - temp_data: Temperature data (independent variable).
        - c_c0_data: Proportion of remaining concentration (dependent variable).
        """

        # Perform curve fitting
        popt, pcov = curve_fit(
            lambda t, A_prime, Ea: self.kinetic_rate_law(
                t, A_prime, Ea, temp_data, self.T_ref),
            time_data,
            c_c0_data,
            p0=self.initial_guess
        )

        # Store optimized parameters
        self.params_opt = popt

        # Calculate predictions for R-squared
        predictions = self.kinetic_rate_law(
            time_data, popt[0], popt[1], temp_data, self.T_ref)
        self.r_squared = r2_score(c_c0_data, predictions)

        # Calculate standard errors of the parameters
        self.param_se = np.sqrt(np.diag(pcov))

        # Calculate 95% confidence intervals
        alpha = 0.05
        n = len(time_data)  # Number of data points
        p = len(popt)  # Number of parameters
        dof = max(0, n - p)  # Degrees of freedom
        t_val = t.ppf(1.0 - alpha / 2.0, dof)

        self.conf_intervals = [
            (param - t_val * se, param + t_val * se)
            for param, se in zip(popt, self.param_se)
        ]

    def print_results(self):
        print(f"Optimized A': {self.params_opt[0]:.5f}")
        print(f"A' Standard Error: {self.param_se[0]:.5f}")
        print(
            f"A' Confidence Interval: {self.conf_intervals[0][0]:.5f}, {self.conf_intervals[0][1]:.5f}")
        print()
        print(f"Optimized E_a: {self.params_opt[1]:.2f}")
        print(f"E_a Standard Error: {self.param_se[1]:.2f}")
        print(
            f"E_a Confidence Interval: {self.conf_intervals[1][0]:.2f}, {self.conf_intervals[1][1]:.2f}")
        print()
        print(f"R-squared: {self.r_squared:.4f}")

    def make_predictions(self, time_data, T):
        return self.kinetic_rate_law(time_data, self.params_opt[0], self.params_opt[1], T, self.T_ref)
