import numpy as np
import matplotlib.pyplot as plt

class LSTM:
    """LSTMのスクラッチ実装（numpy）"""
    def __init__(self, input_size, hidden_size, output_size, learning_rate=0.01):
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.output_size = output_size
        self.lr = learning_rate

        # 結合入力のサイズ
        concat_size = hidden_size + input_size

        # 重みの初期化（Xavier初期化）
        scale = np.sqrt(2.0 / concat_size)

        # 忘却ゲートの重み
        self.Wf = np.random.randn(hidden_size, concat_size) * scale
        self.bf = np.ones((hidden_size, 1))  # 忘却ゲートのバイアスは1で初期化

        # 入力ゲートの重み
        self.Wi = np.random.randn(hidden_size, concat_size) * scale
        self.bi = np.zeros((hidden_size, 1))

        # 候補セル状態の重み
        self.Wc = np.random.randn(hidden_size, concat_size) * scale
        self.bc = np.zeros((hidden_size, 1))

        # 出力ゲートの重み
        self.Wo = np.random.randn(hidden_size, concat_size) * scale
        self.bo = np.zeros((hidden_size, 1))

        # 出力層の重み
        self.Wy = np.random.randn(output_size, hidden_size) * np.sqrt(2.0 / hidden_size)
        self.by = np.zeros((output_size, 1))

    def sigmoid(self, x):
        """数値安定なシグモイド"""
        return np.where(x >= 0,
                       1 / (1 + np.exp(-x)),
                       np.exp(x) / (1 + np.exp(x)))

    def forward(self, inputs):
        """順伝播: inputs は文字インデックスのリスト"""
        T = len(inputs)
        self.T = T
        self.inputs = inputs

        # 各時刻の中間変数を保存
        self.concat_inputs = {}
        self.f_gates = {}
        self.i_gates = {}
        self.c_candidates = {}
        self.o_gates = {}
        self.cell_states = {-1: np.zeros((self.hidden_size, 1))}
        self.hidden_states = {-1: np.zeros((self.hidden_size, 1))}
        self.outputs = {}

        for t in range(T):
            # 入力のone-hotエンコーディング
            x_t = np.zeros((self.input_size, 1))
            x_t[inputs[t]] = 1

            # 結合入力 [h_{t-1}, x_t]
            concat = np.vstack([self.hidden_states[t-1], x_t])
            self.concat_inputs[t] = concat

            # 忘却ゲート
            self.f_gates[t] = self.sigmoid(self.Wf @ concat + self.bf)

            # 入力ゲート
            self.i_gates[t] = self.sigmoid(self.Wi @ concat + self.bi)

            # 候補セル状態
            self.c_candidates[t] = np.tanh(self.Wc @ concat + self.bc)

            # セル状態の更新
            self.cell_states[t] = (self.f_gates[t] * self.cell_states[t-1]
                                    + self.i_gates[t] * self.c_candidates[t])

            # 出力ゲート
            self.o_gates[t] = self.sigmoid(self.Wo @ concat + self.bo)

            # 隠れ状態
            self.hidden_states[t] = self.o_gates[t] * np.tanh(self.cell_states[t])

            # 出力（ソフトマックス）
            logits = self.Wy @ self.hidden_states[t] + self.by
            exp_logits = np.exp(logits - np.max(logits))  # 数値安定性
            self.outputs[t] = exp_logits / np.sum(exp_logits)

        return self.outputs

    def backward(self, targets):
        """逆伝播（BPTT）"""
        T = self.T

        # 勾配の初期化
        dWf = np.zeros_like(self.Wf)
        dWi = np.zeros_like(self.Wi)
        dWc = np.zeros_like(self.Wc)
        dWo = np.zeros_like(self.Wo)
        dbf = np.zeros_like(self.bf)
        dbi = np.zeros_like(self.bi)
        dbc = np.zeros_like(self.bc)
        dbo = np.zeros_like(self.bo)
        dWy = np.zeros_like(self.Wy)
        dby = np.zeros_like(self.by)

        # 次の時刻から伝播してくる勾配
        dh_next = np.zeros((self.hidden_size, 1))
        dc_next = np.zeros((self.hidden_size, 1))

        loss = 0

        for t in reversed(range(T)):
            # 出力層の勾配（交差エントロピー損失）
            dy = self.outputs[t].copy()
            dy[targets[t]] -= 1  # ソフトマックス+交差エントロピーの勾配
            loss -= np.log(self.outputs[t][targets[t]] + 1e-8)

            dWy += dy @ self.hidden_states[t].T
            dby += dy

            # 隠れ状態に対する勾配
            dh = self.Wy.T @ dy + dh_next

            # 出力ゲートの勾配
            tanh_c = np.tanh(self.cell_states[t])
            do = dh * tanh_c
            do_raw = do * self.o_gates[t] * (1 - self.o_gates[t])

            # セル状態に対する勾配
            dc = (dh * self.o_gates[t] * (1 - tanh_c ** 2)) + dc_next

            # 忘却ゲートの勾配
            df = dc * self.cell_states[t-1]
            df_raw = df * self.f_gates[t] * (1 - self.f_gates[t])

            # 入力ゲートの勾配
            di = dc * self.c_candidates[t]
            di_raw = di * self.i_gates[t] * (1 - self.i_gates[t])

            # 候補セル状態の勾配
            dc_cand = dc * self.i_gates[t]
            dc_cand_raw = dc_cand * (1 - self.c_candidates[t] ** 2)

            # 重み・バイアスの勾配を累積
            concat = self.concat_inputs[t]
            dWf += df_raw @ concat.T
            dWi += di_raw @ concat.T
            dWc += dc_cand_raw @ concat.T
            dWo += do_raw @ concat.T
            dbf += df_raw
            dbi += di_raw
            dbc += dc_cand_raw
            dbo += do_raw

            # 前の時刻への勾配
            dconcat = (self.Wf.T @ df_raw + self.Wi.T @ di_raw
                      + self.Wc.T @ dc_cand_raw + self.Wo.T @ do_raw)
            dh_next = dconcat[:self.hidden_size]
            dc_next = dc * self.f_gates[t]

        # 勾配クリッピング
        for grad in [dWf, dWi, dWc, dWo, dbf, dbi, dbc, dbo, dWy, dby]:
            np.clip(grad, -5, 5, out=grad)

        # パラメータ更新
        self.Wf -= self.lr * dWf
        self.Wi -= self.lr * dWi
        self.Wc -= self.lr * dWc
        self.Wo -= self.lr * dWo
        self.bf -= self.lr * dbf
        self.bi -= self.lr * dbi
        self.bc -= self.lr * dbc
        self.bo -= self.lr * dbo
        self.Wy -= self.lr * dWy
        self.by -= self.lr * dby

        return loss.item()

    def generate(self, seed_char, length, char_to_idx, idx_to_char, temperature=1.0):
        """テキスト生成"""
        h = np.zeros((self.hidden_size, 1))
        c = np.zeros((self.hidden_size, 1))
        idx = char_to_idx[seed_char]
        result = [seed_char]

        for _ in range(length):
            x = np.zeros((self.input_size, 1))
            x[idx] = 1
            concat = np.vstack([h, x])

            # 順伝播
            f = self.sigmoid(self.Wf @ concat + self.bf)
            i = self.sigmoid(self.Wi @ concat + self.bi)
            c_cand = np.tanh(self.Wc @ concat + self.bc)
            c = f * c + i * c_cand
            o = self.sigmoid(self.Wo @ concat + self.bo)
            h = o * np.tanh(c)

            # 出力
            logits = self.Wy @ h + self.by
            logits = logits / temperature
            exp_logits = np.exp(logits - np.max(logits))
            probs = exp_logits / np.sum(exp_logits)

            idx = np.random.choice(self.input_size, p=probs.flatten())
            result.append(idx_to_char[idx])

        return ''.join(result)

# --- テキスト生成タスク ---

# サンプルテキスト（学習データ）
text = """the quick brown fox jumps over the lazy dog
a stitch in time saves nine
all that glitters is not gold
actions speak louder than words
knowledge is power and power is knowledge
the pen is mightier than the sword
fortune favors the bold and brave
to be or not to be that is the question
early bird catches the worm every morning
practice makes perfect in every endeavor"""

# 文字のインデックス化
chars = sorted(set(text))
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}
vocab_size = len(chars)

print(f"語彙サイズ: {vocab_size}")
print(f"テキスト長: {len(text)}")

# 学習データの準備
data = [char_to_idx[ch] for ch in text]

# LSTMモデルの構築
hidden_size = 64
lstm = LSTM(input_size=vocab_size, hidden_size=hidden_size,
            output_size=vocab_size, learning_rate=0.005)

# 学習
n_epochs = 200
seq_length = 30  # 1回の学習で使う系列長
losses = []

for epoch in range(n_epochs):
    epoch_loss = 0
    n_batches = 0

    # テキストをseq_length単位で処理
    for start in range(0, len(data) - seq_length - 1, seq_length):
        inputs = data[start:start + seq_length]
        targets = data[start + 1:start + seq_length + 1]

        # 順伝播
        lstm.forward(inputs)

        # 逆伝播
        loss = lstm.backward(targets)
        epoch_loss += loss
        n_batches += 1

    avg_loss = epoch_loss / n_batches
    losses.append(avg_loss)

    if (epoch + 1) % 40 == 0:
        print(f"\nEpoch {epoch+1}/{n_epochs}, Loss: {avg_loss:.4f}")
        generated = lstm.generate('t', 80, char_to_idx, idx_to_char, temperature=0.8)
        print(f"Generated: {generated}")

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(16, 5))

# 学習曲線
axes[0].plot(losses, linewidth=2)
axes[0].set_xlabel('Epoch', fontsize=12)
axes[0].set_ylabel('Average Loss', fontsize=12)
axes[0].set_title('LSTM Training Loss (Character-level Language Model)', fontsize=14)
axes[0].grid(True, alpha=0.3)

# ゲートの値の可視化（最後のエピソードの忘却ゲート値）
sample_input = data[:seq_length]
lstm.forward(sample_input)

f_gate_values = np.array([lstm.f_gates[t].flatten()[:8] for t in range(seq_length)])
im = axes[1].imshow(f_gate_values.T, aspect='auto', cmap='RdYlGn', vmin=0, vmax=1)
axes[1].set_xlabel('Time Step', fontsize=12)
axes[1].set_ylabel('Hidden Unit', fontsize=12)
axes[1].set_title('Forget Gate Values (First 8 Units)', fontsize=14)
plt.colorbar(im, ax=axes[1])

# 入力文字をx軸ラベルとして表示
chars_label = [idx_to_char[idx] for idx in sample_input]
axes[1].set_xticks(range(0, seq_length, 5))
axes[1].set_xticklabels([chars_label[i] for i in range(0, seq_length, 5)])

plt.tight_layout()
plt.show()

# 最終的なテキスト生成
print("\n--- 学習済みLSTMによるテキスト生成 ---")
for seed in ['t', 'a', 'k']:
    generated = lstm.generate(seed, 100, char_to_idx, idx_to_char, temperature=0.7)
    print(f"\nSeed='{seed}': {generated}")

import numpy as np
import matplotlib.pyplot as plt

def compute_gradient_flow_rnn(hidden_size, seq_len, n_trials=50):
    """通常のRNNの勾配流を計算"""
    norms = np.zeros(seq_len)
    for _ in range(n_trials):
        # ランダムな重み行列
        Wh = np.random.randn(hidden_size, hidden_size) * 0.5 / np.sqrt(hidden_size)
        # ランダムな隠れ状態
        h = np.random.randn(hidden_size, 1) * 0.1

        jacobian_prod = np.eye(hidden_size)
        for t in range(seq_len):
            # h_t = tanh(Wh @ h_{t-1} + ...)
            a = Wh @ h
            h = np.tanh(a)
            # ヤコビ行列: diag(1 - h^2) @ Wh
            J = np.diag((1 - h ** 2).flatten()) @ Wh
            jacobian_prod = J @ jacobian_prod
            norms[t] += np.linalg.norm(jacobian_prod)

    return norms / n_trials

def compute_gradient_flow_lstm(hidden_size, seq_len, n_trials=50):
    """LSTMの勾配流を計算（セル状態経由）"""
    norms = np.zeros(seq_len)
    for _ in range(n_trials):
        # 忘却ゲートの値（0.9前後に設定、実際の学習済みLSTMを想定）
        forget_gates = 0.85 + 0.1 * np.random.rand(seq_len, hidden_size)

        gradient_prod = np.ones(hidden_size)
        for t in range(seq_len):
            # ∂C_t/∂C_{t-1} ≈ diag(f_t)
            gradient_prod *= forget_gates[t]
            norms[t] += np.linalg.norm(gradient_prod)

    return norms / n_trials

# 実験の実行
hidden_size = 32
seq_len = 100

rnn_norms = compute_gradient_flow_rnn(hidden_size, seq_len)
lstm_norms = compute_gradient_flow_lstm(hidden_size, seq_len)

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# 線形スケール
axes[0].plot(range(1, seq_len+1), rnn_norms, label='Simple RNN', linewidth=2)
axes[0].plot(range(1, seq_len+1), lstm_norms, label='LSTM (Cell State)', linewidth=2)
axes[0].set_xlabel('Time Steps Back', fontsize=12)
axes[0].set_ylabel('Gradient Norm', fontsize=12)
axes[0].set_title('Gradient Flow: RNN vs LSTM (Linear Scale)', fontsize=14)
axes[0].legend(fontsize=12)
axes[0].grid(True, alpha=0.3)

# 対数スケール
axes[1].semilogy(range(1, seq_len+1), rnn_norms + 1e-20, label='Simple RNN', linewidth=2)
axes[1].semilogy(range(1, seq_len+1), lstm_norms + 1e-20, label='LSTM (Cell State)', linewidth=2)
axes[1].set_xlabel('Time Steps Back', fontsize=12)
axes[1].set_ylabel('Gradient Norm (log scale)', fontsize=12)
axes[1].set_title('Gradient Flow: RNN vs LSTM (Log Scale)', fontsize=14)
axes[1].legend(fontsize=12)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 数値結果
print(f"100ステップ後の勾配ノルム:")
print(f"  RNN:  {rnn_norms[-1]:.2e}")
print(f"  LSTM: {lstm_norms[-1]:.2e}")
print(f"  比率: {lstm_norms[-1] / (rnn_norms[-1] + 1e-30):.2e}")