"""
Temperature Profile Generator Module

This module provides various methods for generating temperature profiles for atmospheric modeling.
It supports multiple profile types including isothermal, Guillot, Madhu, personalized linear, and
four-node spline interpolation profiles.

The module includes uncertainty propagation capabilities for all profile types, allowing
for error estimation through parameter perturbation.
"""

import numpy as np
from numpy import interp
from scipy.special import comb
from petitRADTRANS.physics import temperature_profile_function_guillot_global
import pyratbay.atmosphere as pa


def bernstein_poly(i: int, n: int, t: np.ndarray) -> np.ndarray:
    """
    Calculate the Bernstein polynomial of degree n and index i.

    Args:
        i (int): Index of the Bernstein polynomial
        n (int): Degree of the polynomial
        t (np.ndarray): Array of points at which to evaluate the polynomial

    Returns:
        np.ndarray: Values of the Bernstein polynomial at points t
    """
    return comb(n, i) * (t ** (n - i)) * (1 - t) ** i


def bezier_curve(points, nTimes: int = 1000) -> tuple[np.ndarray, np.ndarray]:
    """
    Generate a Bezier curve from a set of control points.

    Args:
        points (list): List of [x,y] control points defining the curve
        nTimes (int, optional): Number of points to generate along the curve. Defaults to 1000.

    Returns:
        tuple[np.ndarray, np.ndarray]: Arrays of x and y coordinates along the curve
    """
    nPoints = len(points)
    xPoints = np.array([p[0] for p in points])
    yPoints = np.array([p[1] for p in points])
    t = np.linspace(0.0, 1.0, nTimes)
    polynomial_array = np.array([bernstein_poly(i, nPoints - 1, t) for i in range(0, nPoints)])
    xvals = np.dot(xPoints, polynomial_array)
    yvals = np.dot(yPoints, polynomial_array)
    return xvals, yvals


class TemperatureProfile:
    """
    A class for generating various types of temperature profiles for atmospheric modeling.

    This class supports multiple profile types including isothermal, Guillot, Madhu,
    personalized linear, and four-node spline interpolation profiles. It includes
    built-in uncertainty propagation capabilities.

    Attributes:
        pressures (np.ndarray): Array of pressure levels
        parameters (dict): Dictionary of parameter objects with values and errors
        gravity (float): Gravitational acceleration
        error (bool): Whether to compute error profiles
        rng (np.random.Generator): Random number generator for error propagation
    """

    def __init__(self, pressures: np.ndarray, parameters: dict, gravity: float, 
                 error: bool = False, rng: np.random.Generator = None):
        """
        Initialize the TemperatureProfile generator.

        Args:
            pressures (np.ndarray): Array of pressure levels
            parameters (dict): Dictionary of parameter objects with values and errors
            gravity (float): Gravitational acceleration
            error (bool, optional): Whether to compute error profiles. Defaults to False.
            rng (np.random.Generator, optional): Random number generator. If None, creates a new one.
        """
        self.pressures = np.asarray(pressures, dtype=float)
        self.parameters = parameters
        self.gravity = float(gravity)
        self.error = bool(error)
        self.rng = rng if rng is not None else np.random.default_rng()

    def return_rnd_for_error(self, name: str, amount: int = 100) -> np.ndarray:
        """
        Generate random samples for a parameter within its error bounds.

        Args:
            name (str): Name of the parameter to sample
            amount (int, optional): Number of samples to generate. Defaults to 100.

        Returns:
            np.ndarray: Array of random samples within parameter's error bounds

        Raises:
            KeyError: If parameter name is not found in self.parameters
        """
        if name not in self.parameters:
            raise KeyError(f"Parameter '{name}' not found in parameters dictionary")
            
        param = self.parameters[name]
        return self.rng.uniform(
            param.value - param.error,
            param.value + param.error,
            amount
        )

    def isot(self) -> tuple[np.ndarray, None]:
        """
        Compute an isothermal temperature profile.

        The isothermal model returns a constant temperature for all pressure levels,
        with no associated uncertainty.

        Returns:
            tuple[np.ndarray, None]: 
                - temperature: Array of the isothermal temperature (T0) at each pressure
                - None: No error profile is computed for the isothermal model

        Raises:
            KeyError: If T0 parameter is not found in parameters dictionary
        """
        if "T0" not in self.parameters:
            raise KeyError("T0 parameter not found in parameters dictionary")
            
        T0 = self.parameters["T0"].value
        temperature = T0 * np.ones_like(self.pressures)
        return temperature, None

    def guillot(self) -> tuple[np.ndarray, np.ndarray | None]:
        """
        Compute the temperature profile using the Guillot model.

        This method computes the nominal temperature profile using the Guillot model.
        If error propagation is enabled (self.error is True), it also computes an ensemble
        of temperature profiles using random perturbations of the input parameters.

        Returns:
            tuple[np.ndarray, np.ndarray | None]:
                - temperature: Nominal temperature profile
                - error_temperatures: If self.error is True, array of shape (len(pressures), num_samples)
                  containing perturbed profiles. Otherwise None.

        Raises:
            KeyError: If required parameters (kappa_IR, gamma_g, T_int, T0) are not found
        """
        required_params = ["kappa_IR", "gamma_g", "T_int", "T0"]
        for param in required_params:
            if param not in self.parameters:
                raise KeyError(f"Required parameter '{param}' not found in parameters dictionary")

        # Extract nominal parameter values
        kappa_IR = self.parameters["kappa_IR"].value
        gamma_g = self.parameters["gamma_g"].value
        T_int = self.parameters["T_int"].value
        T0 = self.parameters["T0"].value

        # Compute the nominal temperature profile
        temperature = temperature_profile_function_guillot_global(
            self.pressures, kappa_IR, gamma_g, self.gravity, T_int, T0
        )

        error_temperatures = None
        if self.error:
            num_samples = 100
            error_temperatures = np.zeros((len(self.pressures), num_samples))
            
            # Generate random samples for parameters
            ka = 10**self.return_rnd_for_error("kappa_IR")
            ga = 10**self.return_rnd_for_error("gamma_g")
            Ti_err = self.return_rnd_for_error("T_int")
            te = self.return_rnd_for_error("T0")

            # Compute the temperature profile for each set of perturbed parameters
            for i in range(num_samples):
                error_temperatures[:, i] = temperature_profile_function_guillot_global(
                    self.pressures, ka[i], ga[i], self.gravity, Ti_err[i], te[i]
                )

        return temperature, error_temperatures

    def madhu(self) -> tuple[np.ndarray, np.ndarray | None]:
        """
        Compute the temperature profile using the Madhu model.

        This method calculates the nominal temperature profile using the Madhu model
        with the following parameters:
        - p1, p2, p3: Pressure control points
        - alpha1, alpha2: Temperature gradient parameters
        - T0: Base temperature

        If error propagation is enabled, it computes an ensemble of profiles by
        sampling parameters within their error bounds.

        Returns:
            tuple[np.ndarray, np.ndarray | None]:
                - temperature: Nominal temperature profile
                - error_temperatures: If self.error is True, array of shape (len(pressures), 100)
                  containing perturbed profiles. Otherwise None.

        Raises:
            KeyError: If required parameters are not found in parameters dictionary
        """
        required_params = ["p1", "p2", "p3", "alpha1", "alpha2", "T0"]
        for param in required_params:
            if param not in self.parameters:
                raise KeyError(f"Required parameter '{param}' not found in parameters dictionary")

        error_temperatures = None
        tp_madhu = pa.tmodels.Madhu(self.pressures)

        # Compute the nominal temperature profile with central parameter values
        params = [
            self.parameters["p1"].value,
            self.parameters["p2"].value,
            self.parameters["p3"].value,
            self.parameters["alpha1"].value,
            self.parameters["alpha2"].value,
            self.parameters["T0"].value
        ]
        temperature = tp_madhu(params)

        if self.error:
            num_samples = 100
            error_temperatures = np.zeros((len(self.pressures), num_samples))
            
            # Generate random samples for each parameter
            p1_err = self.return_rnd_for_error("p1")
            p2_err = self.return_rnd_for_error("p2")
            p3_err = self.return_rnd_for_error("p3")
            alpha1_err = self.return_rnd_for_error("alpha1")
            alpha2_err = self.return_rnd_for_error("alpha2")
            T0_err = self.return_rnd_for_error("T0")

            # Compute temperature profiles for perturbed parameters
            for i in range(num_samples):
                params_err = [
                    p1_err[i], p2_err[i], p3_err[i],
                    alpha1_err[i], alpha2_err[i], T0_err[i]
                ]
                error_temperatures[:, i] = tp_madhu(params_err)

        return temperature, error_temperatures

    def personalized(self) -> tuple[np.ndarray, np.ndarray | None]:
        """
        Compute a personalized temperature profile based on pressure thresholds.

        The profile is computed as a linear function of log10(pressure) between two
        pressure thresholds (P_low and P_high) with corresponding temperatures
        (T_low and T_high). The temperature is:
        - T_low for pressures > 10^P_low
        - T_high for pressures < 10^P_high
        - Linearly interpolated in between

        Returns:
            tuple[np.ndarray, np.ndarray | None]:
                - temperature: Nominal temperature profile
                - error_temperatures: If self.error is True, array of shape (len(pressures), 100)
                  containing perturbed profiles. Otherwise None.

        Raises:
            KeyError: If required parameters are not found in parameters dictionary
        """
        required_params = ["P_low", "P_high", "T_low", "T_high"]
        for param in required_params:
            if param not in self.parameters:
                raise KeyError(f"Required parameter '{param}' not found in parameters dictionary")

        # Retrieve central parameter values
        p1 = self.parameters["P_low"].value
        p2 = self.parameters["P_high"].value
        t1 = self.parameters["T_low"].value
        t2 = self.parameters["T_high"].value

        # Compute the nominal temperature profile
        factor = (t1 - t2) / (p1 - p2)
        log_pressures = np.log10(self.pressures)
        temperature = log_pressures * factor + t1 - factor * p1

        # Apply piecewise limits
        temperature[self.pressures > 10**p1] = t1
        temperature[self.pressures < 10**p2] = t2

        error_temperatures = None
        if self.error:
            num_samples = 100
            error_temperatures = np.zeros((len(self.pressures), num_samples))
            
            # Generate random samples for each parameter
            ktl = self.return_rnd_for_error("T_low")
            kth = self.return_rnd_for_error("T_high")
            kpl = self.return_rnd_for_error("P_low")
            kph = self.return_rnd_for_error("P_high")

            for i in range(num_samples):
                # Initialize with T_low
                tempio = t1 * np.ones_like(self.pressures)
                
                # Apply thresholds
                tempio[self.pressures <= 10**kph[i]] = kth[i]
                tempio[self.pressures >= 10**kpl[i]] = ktl[i]
                
                # Interpolate in between
                inside = (self.pressures < 10**kpl[i]) & (self.pressures > 10**kph[i])
                if np.any(inside):
                    tempio[inside] = np.linspace(kth[i], ktl[i], np.sum(inside))
                
                error_temperatures[:, i] = tempio

        return temperature, error_temperatures

    def FourNodeSpline(self) -> tuple[np.ndarray, np.ndarray | None]:
        """
        Compute the temperature profile using a four-node Bezier spline interpolation.

        The profile is created by interpolating between four node points using a Bezier curve:
        - Temperature nodes: [T3_node, T2_node, T1_node, T0_node]
        - Pressure nodes: [log10(max(pressure)), P2_node, P1_node, log10(min(pressure))]

        The interpolation is performed in log-pressure space.

        Returns:
            tuple[np.ndarray, np.ndarray | None]:
                - temperature: Nominal temperature profile
                - error_temperatures: If self.error is True, array of shape (len(pressures), 100)
                  containing perturbed profiles. Otherwise None.

        Raises:
            KeyError: If required parameters are not found in parameters dictionary
        """
        required_params = ["T0_node", "T1_node", "T2_node", "T3_node", "P1_node", "P2_node"]
        for param in required_params:
            if param not in self.parameters:
                raise KeyError(f"Required parameter '{param}' not found in parameters dictionary")

        # Retrieve central node values
        T0_node = self.parameters["T0_node"].value
        T1_node = self.parameters["T1_node"].value
        T2_node = self.parameters["T2_node"].value
        T3_node = self.parameters["T3_node"].value
        P1_node = self.parameters["P1_node"].value
        P2_node = self.parameters["P2_node"].value

        # Compute the logarithm of the pressures
        logP = np.log10(self.pressures)
        nlev = logP.shape[0]

        # Define the node positions
        pressure_nodes = [logP.max(), P2_node, P1_node, logP.min()]
        temperature_nodes = [T3_node, T2_node, T1_node, T0_node]
        pts = np.column_stack((temperature_nodes, pressure_nodes))

        # Generate the Bezier curve
        t, p = bezier_curve(pts, nlev)
        temperature = interp(logP, p, t)

        error_temperatures = None
        if self.error:
            n_samples = 100
            error_temperatures = np.zeros((len(self.pressures), n_samples))
            
            # Generate random samples for each node parameter
            errt0 = self.return_rnd_for_error("T0_node")
            errt1 = self.return_rnd_for_error("T1_node")
            errt2 = self.return_rnd_for_error("T2_node")
            errt3 = self.return_rnd_for_error("T3_node")
            errp1 = self.return_rnd_for_error("P1_node")
            errp2 = self.return_rnd_for_error("P2_node")

            for i in range(n_samples):
                # Create perturbed node values
                pressure_nodes_sample = [logP.max(), errp2[i], errp1[i], logP.min()]
                temperature_nodes_sample = [errt3[i], errt2[i], errt1[i], errt0[i]]
                pts_sample = np.column_stack((temperature_nodes_sample, pressure_nodes_sample))
                
                # Generate perturbed Bezier curve
                t_sample, p_sample = bezier_curve(pts_sample, nlev)
                error_temperatures[:, i] = interp(logP, p_sample, t_sample)

        return temperature, error_temperatures
