"""
This module provides the main EchoStateNetwork instance, which can be trained to forecast time series
"""


from dataclasses import dataclass, field
from typing import Optional, Callable, Union
import warnings

import numpy as np

from .time_series import TimeSeries


class InstabilityWarning(Warning):

    """
    Warning indicating some sort of numerical instability
    """

    pass


@dataclass
class EchoStateNetwork:

    """
    ESN class for making prediction
    """

    reservoir_dimensionality: int
    """dimensionality of reservoir, should be much larger than input dimension"""
    input_dimensionality: int
    """expected dimensionality of input time series"""
    w_in: Optional[np.typing.NDArray] = None
    """matrix defining input -> reservoir mapping, autogenerated if not specified"""
    regularization_parameter: float = 1.0e-10
    """regularization parameter for fit, should be small, defaults to 1.0e-10"""
    generator: Optional[np.random.Generator] = None
    """random generator, auto generated if not specified"""
    activation_function: Callable[[np.typing.NDArray], np.typing.NDArray] = np.tanh
    """activation function squishing reservoir states, defaults to hyperbolic tangent"""
    w_out: np.typing.NDArray = field(init=False)
    """matrix defining reservoir -> output mapping; trainable and should not be set at initialization"""
    feedback_loop_guess: Union[float, np.typing.NDArray] = field(init=False)
    """starting warmup for prediction; trainable and should not be set at initialization"""
    timestep: float = field(init=False)
    """detected timestep of inputted TimeSeries; trainable and should not be set at initialization"""
    final_time: float = field(init=False)
    """final time of inputted TimeSeries; trainable and should not be set at initialization"""

    def __post_init__(self):

        """
        Need to initialize a generator and a W_in.
        The generator attribute here is responsible for random sampling
        """

        if not self.generator:
            self.generator = np.random.default_rng(seed=0)

        if self.w_in is None:
            self.w_in = self.generator.uniform(
                low=-0.5,
                high=0.5,
                size=(self.reservoir_dimensionality, self.input_dimensionality)
            )

    def train(self, input_time_series: TimeSeries, pinv: bool = False, max_condition_number: float = 10.0):

        """
        Training step, solving argmin_W ||Y - WX||^2 + λ||W||^2

        Args:
            input_time_series (TimeSeries):
                The input TimeSeries instance to train on

            pinv (bool):
                Uses the Moore-Penrose inverse in the training if true. Defaults to false.

            max_condition_number (float):
                Maximum condition number of the training data.
                A larger condition number will induce an InstabilityWarning
        """

        # assert that all timesteps are constant
        self.timestep = input_time_series.timestep
        self.final_time = input_time_series.times[-1]

        time_series_array = input_time_series.dependent_variable
        independent_variables = self.activation_function(self.w_in @ time_series_array[:-1, :].T).T
        dependent_variables = time_series_array[1:]

        regularization = self.regularization_parameter * np.eye(independent_variables.shape[1])

        if pinv:
            w_out_transpose = np.linalg.pinv(independent_variables) @ dependent_variables
        else:
            x = independent_variables.T @ independent_variables + regularization

            # conditional number 
            cond_num = np.linalg.cond(x)
            if cond_num > max_condition_number:
                warnings.warn(
                    f"Condition Number {cond_num:.2E} is greater than {max_condition_number}. "
                    f"consider passing pinv = True in {self.__class__.__name__}.train() "
                    f"or increasing {self.__class__.__name__}.regularization_parameter",
                    InstabilityWarning
                )

            w_out_transpose = np.linalg.solve(x, independent_variables.T @ dependent_variables)
            
        self.w_out = w_out_transpose.T
        self.feedback_loop_guess = time_series_array[-1]

    def predict(self, horizon: int) -> TimeSeries:

        """
        Prediction step, using the previous training to forecast into the future

        Args:
            horizon (int): The number of steps to predict into the future

        Returns:
            TimeSeries: The resulting predicted time series
        """

        if self.w_out is None or self.w_in is None:
            raise ValueError("need to train before predicting")

        # initialize predictions and reservoir states to populate later
        predictions = np.zeros((horizon, self.input_dimensionality))

        # perform feedback loop
        for i in range(horizon):
            if i == 0:
                predictions[i, :] = self.w_out @ self.activation_function(self.w_in @ self.feedback_loop_guess)
                continue
            predictions[i, :] = self.w_out @ self.activation_function(self.w_in @ predictions[i - 1, :])

        return TimeSeries(
            dependent_variable=predictions,
            times=self.final_time + self.timestep + np.linspace(0, horizon * self.timestep, horizon)
        )
