import numpy as np
import matplotlib.pyplot as plt

# --- リザバー層のクラス定義 ---
class Reservoir:
    def __init__(self, n_neurons=100, spectral_radius=0.9, seed=42):
        """
        n_neurons: リザバーのニューロン数（この数だけ特徴量ができる）
        spectral_radius: スペクトル半径（リザバーの「記憶」の長さに影響）
        """
        self.n_neurons = n_neurons
        self.rng = np.random.RandomState(seed)
        
        # 1. 内部結合行列 W_res (ニューロン同士の結合)
        #    ランダムに生成し、スペクトル半径でスケーリングします
        W = self.rng.normal(0, 1, (n_neurons, n_neurons))
        rho_W = np.max(np.abs(np.linalg.eigvals(W)))
        self.W_res = W * (spectral_radius / rho_W)
        
        # 2. 入力結合行列 W_in (入力信号 -> ニューロン)
        #    入力は1次元と仮定
        self.W_in = self.rng.normal(0, 1, (n_neurons, 1))
        
        # 状態ベクトルの初期化
        self.state = np.zeros((n_neurons, 1))

    def forward(self, input_val):
        """
        1ステップ分の入力(スカラー)を受け取り、新しい状態を返す
        x(t) = tanh( W_in * u(t) + W_res * x(t-1) )
        """
        # 入力 u(t)
        u = np.array([[input_val]])
        
        # 更新式
        pre_activation = np.dot(self.W_in, u) + np.dot(self.W_res, self.state)
        self.state = np.tanh(pre_activation)
        
        # 1次元配列にして返す
        return self.state.flatten()

# --- メイン処理 ---
def main():
    # 1. データの読み込み
    filename = "mackey_glass.txt"
    try:
        data = np.loadtxt(filename)
        print(f"データ読み込み完了: {len(data)} ステップ")
    except FileNotFoundError:
        print(f"エラー: {filename} が見つかりません。gen_data.py を実行しましたか？")
        return

    # 2. リザバーの準備
    n_neurons = 50  # 実験用に50個とします（本番は100~500程度）
    reservoir = Reservoir(n_neurons=n_neurons)

    # 3. データを流し込んで「状態行列」を作る
    states = []
    
    # ウォッシュアウト（初期過渡応答を捨てる期間）
    washout = 50
    
    print("リザバー変換を実行中...")
    for t, u in enumerate(data):
        s = reservoir.forward(u)
        if t >= washout:
            states.append(s)
            
    # NumPy配列に変換 ( サンプル数 x ニューロン数 )
    X = np.array(states)
    
    # 4. 「正解データ」を作る
    #    タスク: 「今の入力」から「次のステップの入力」を予測する
    #    X[t] に対応する正解は data[t + washout + 1] です
    #    (最後の1点は次の正解がないので使えません)
    
    # X は data[washout] から data[end] までの状態
    # y は data[washout+1] から data[end] までの値
    
    # サイズ合わせ
    X_train = X[:-1]      # 最後の1行を捨てる
    y_train = data[washout+1:] # 最初のwashout+1個を捨てる
    
    print("-" * 30)
    print(f"変換完了:")
    print(f"  入力データ数    : {len(data)}")
    print(f"  リザバー状態(X) : {X_train.shape}  (サンプル数, 変数/ニューロン数)")
    print(f"  正解データ(y)   : {y_train.shape}  (サンプル数, )")
    print("-" * 30)
    print("この X と y を使って e04gnf で『X * w ≒ y』となる w を求めます。")

    # 5. 可視化（リザバーの中で何が起きているか確認）
    plt.figure(figsize=(12, 6))
    
    # 元データ
    plt.subplot(2, 1, 1)
    plt.plot(y_train[:200], 'k', label="Target Signal (Original)")
    plt.title("Original Signal vs Reservoir Neuron Activations")
    plt.legend(loc="upper right")
    plt.grid(True)
    
    # ニューロンの反応（最初の5個だけ表示）
    plt.subplot(2, 1, 2)
    plt.plot(X_train[:200, :5]) 
    plt.title("Activations of first 5 neurons")
    plt.ylabel("State value")
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()