import reservoirpy as rpy
import torch
from numpy.linalg import eigvals
from reservoirpy.nodes import Reservoir, Ridge
import numpy as np
from dataclasses import dataclass


@dataclass
class KoopmanESNConfig:
    units: int
    lr: float
    sr: float
    sp: float
    ridge: float
    train_start: int
    train_len: int
    train_warmup: int
    predict_warmup: int
    predict_len: int
    state_dim: int


class KoopmanESN:
    def __init__(self, A, config: KoopmanESNConfig):
        super().__init__()

        self.esn_model = None
        self.config = config
        self.A = A/100

    def W_initial(self, A):
        W_koop = A
        W = W_koop

        # 5. 调整谱半径
        current_radius = np.max(np.abs(eigvals(W)))
        W = (self.config.sr / current_radius) * W

        return W

    def koopmanesn_train(self, data_train):
        rpy.verbosity(0)
        rpy.set_seed(42)

        W = self.W_initial(self.A)
        Win = np.eye(self.config.units, self.config.units)
        reservoir = Reservoir(
            units=self.config.units,
            lr=self.config.lr,
            sr=self.config.sr,
            W=W,
            input_scaling=1.0,
            Win=Win,
            activation='identity'
        )

        readout = Ridge(ridge=self.config.ridge)

        esn_model = reservoir >> readout

        data_train_in = data_train[:-1, :]
        data_train_out = data_train[1:, :]

        self.esn_model = esn_model.fit(
            data_train_in,
            data_train_out,
            warmup=self.config.train_warmup
        )

    def predict(self, data_warm):
        warmup_data = self.esn_model.run(data_warm, reset=True)
        x = warmup_data[-1, :].reshape(1, -1)
        data_pre = np.empty((self.config.predict_len, self.config.state_dim))
        for i in range(self.config.predict_len):
            data_pre[i, :] = x
            x = self.esn_model(x)
        return data_pre