import numpy as np
import matplotlib.pyplot as plt

# ----------------------------
# シグモイド関数とtanh関数
# ----------------------------
def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

def tanh(x):
    return np.tanh(x)

# ----------------------------
# GRUクラス（スクラッチ実装）
# ----------------------------
class GRU:
    """GRU（Gated Recurrent Unit）のスクラッチ実装"""

    def __init__(self, input_dim, hidden_dim, output_dim, lr=0.01):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.lr = lr

        # Xavier初期化
        scale_ih = np.sqrt(2.0 / (input_dim + hidden_dim))
        scale_hh = np.sqrt(2.0 / (hidden_dim + hidden_dim))
        scale_ho = np.sqrt(2.0 / (hidden_dim + output_dim))

        # 更新ゲートのパラメータ
        self.Wz = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Uz = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.bz = np.zeros((hidden_dim, 1))

        # リセットゲートのパラメータ
        self.Wr = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Ur = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.br = np.zeros((hidden_dim, 1))

        # 候補隠れ状態のパラメータ
        self.Wh = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Uh = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.bh = np.zeros((hidden_dim, 1))

        # 出力層のパラメータ
        self.Wy = np.random.randn(output_dim, hidden_dim) * scale_ho
        self.by = np.zeros((output_dim, 1))

    def forward(self, xs):
        """順伝播: xs は (T, input_dim, 1) のリストまたは配列"""
        T = len(xs)
        self.xs = xs
        self.zs = {}   # 更新ゲート
        self.rs = {}   # リセットゲート
        self.h_tildes = {}  # 候補隠れ状態
        self.hs = {-1: np.zeros((self.hidden_dim, 1))}  # 隠れ状態
        self.ys = {}   # 出力

        for t in range(T):
            x_t = xs[t]
            h_prev = self.hs[t - 1]

            # 更新ゲート
            self.zs[t] = sigmoid(self.Wz @ x_t + self.Uz @ h_prev + self.bz)

            # リセットゲート
            self.rs[t] = sigmoid(self.Wr @ x_t + self.Ur @ h_prev + self.br)

            # 候補隠れ状態
            self.h_tildes[t] = tanh(
                self.Wh @ x_t + self.Uh @ (self.rs[t] * h_prev) + self.bh
            )

            # 隠れ状態の更新
            self.hs[t] = self.zs[t] * h_prev + (1 - self.zs[t]) * self.h_tildes[t]

            # 出力
            self.ys[t] = self.Wy @ self.hs[t] + self.by

        return self.ys

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

        # パラメータ勾配の初期化
        dWz = np.zeros_like(self.Wz)
        dUz = np.zeros_like(self.Uz)
        dbz = np.zeros_like(self.bz)
        dWr = np.zeros_like(self.Wr)
        dUr = np.zeros_like(self.Ur)
        dbr = np.zeros_like(self.br)
        dWh = np.zeros_like(self.Wh)
        dUh = np.zeros_like(self.Uh)
        dbh = np.zeros_like(self.bh)
        dWy = np.zeros_like(self.Wy)
        dby = np.zeros_like(self.by)

        dh_next = np.zeros((self.hidden_dim, 1))
        loss = 0.0

        for t in reversed(range(T)):
            # 出力層の勾配（MSE損失）
            dy = self.ys[t] - targets[t]
            loss += 0.5 * np.sum(dy ** 2)

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

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

            # 更新ゲートの勾配
            dz = dh * (self.hs[t - 1] - self.h_tildes[t])
            dz_raw = dz * self.zs[t] * (1 - self.zs[t])

            # 候補隠れ状態の勾配
            dh_tilde = dh * (1 - self.zs[t])
            dh_tilde_raw = dh_tilde * (1 - self.h_tildes[t] ** 2)

            # リセットゲートの勾配
            dr = (self.Uh.T @ dh_tilde_raw) * self.hs[t - 1]
            dr_raw = dr * self.rs[t] * (1 - self.rs[t])

            # パラメータ勾配の累積
            dWz += dz_raw @ self.xs[t].T
            dUz += dz_raw @ self.hs[t - 1].T
            dbz += dz_raw

            dWr += dr_raw @ self.xs[t].T
            dUr += dr_raw @ self.hs[t - 1].T
            dbr += dr_raw

            dWh += dh_tilde_raw @ self.xs[t].T
            dUh += dh_tilde_raw @ (self.rs[t] * self.hs[t - 1]).T
            dbh += dh_tilde_raw

            # 前時刻への隠れ状態勾配の伝播
            dh_next = (dh * self.zs[t]                        # 直接経路
                      + self.Uz.T @ dz_raw                    # 更新ゲート経由
                      + self.Ur.T @ dr_raw                    # リセットゲート経由
                      + (self.Uh.T @ dh_tilde_raw) * self.rs[t])  # 候補経由

        # 勾配クリッピング
        for grad in [dWz, dUz, dbz, dWr, dUr, dbr, dWh, dUh, dbh, dWy, dby]:
            np.clip(grad, -5, 5, out=grad)

        # パラメータ更新
        self.Wz -= self.lr * dWz
        self.Uz -= self.lr * dUz
        self.bz -= self.lr * dbz
        self.Wr -= self.lr * dWr
        self.Ur -= self.lr * dUr
        self.br -= self.lr * dbr
        self.Wh -= self.lr * dWh
        self.Uh -= self.lr * dUh
        self.bh -= self.lr * dbh
        self.Wy -= self.lr * dWy
        self.by -= self.lr * dby

        return loss

class LSTM:
    """LSTM（Long Short-Term Memory）のスクラッチ実装（比較用）"""

    def __init__(self, input_dim, hidden_dim, output_dim, lr=0.01):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.lr = lr

        scale_ih = np.sqrt(2.0 / (input_dim + hidden_dim))
        scale_hh = np.sqrt(2.0 / (hidden_dim + hidden_dim))
        scale_ho = np.sqrt(2.0 / (hidden_dim + output_dim))

        # 忘却ゲート
        self.Wf = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Uf = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.bf = np.ones((hidden_dim, 1))  # 忘却ゲートのバイアスは1で初期化

        # 入力ゲート
        self.Wi = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Ui = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.bi = np.zeros((hidden_dim, 1))

        # 出力ゲート
        self.Wo = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Uo = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.bo = np.zeros((hidden_dim, 1))

        # 候補セル
        self.Wc = np.random.randn(hidden_dim, input_dim) * scale_ih
        self.Uc = np.random.randn(hidden_dim, hidden_dim) * scale_hh
        self.bc = np.zeros((hidden_dim, 1))

        # 出力層
        self.Wy = np.random.randn(output_dim, hidden_dim) * scale_ho
        self.by = np.zeros((output_dim, 1))

    def forward(self, xs):
        T = len(xs)
        self.xs = xs
        self.fs, self.ins, self.os, self.c_tildes = {}, {}, {}, {}
        self.cs = {-1: np.zeros((self.hidden_dim, 1))}
        self.hs = {-1: np.zeros((self.hidden_dim, 1))}
        self.ys = {}

        for t in range(T):
            x_t = xs[t]
            h_prev = self.hs[t - 1]
            c_prev = self.cs[t - 1]

            self.fs[t] = sigmoid(self.Wf @ x_t + self.Uf @ h_prev + self.bf)
            self.ins[t] = sigmoid(self.Wi @ x_t + self.Ui @ h_prev + self.bi)
            self.os[t] = sigmoid(self.Wo @ x_t + self.Uo @ h_prev + self.bo)
            self.c_tildes[t] = tanh(self.Wc @ x_t + self.Uc @ h_prev + self.bc)

            self.cs[t] = self.fs[t] * c_prev + self.ins[t] * self.c_tildes[t]
            self.hs[t] = self.os[t] * tanh(self.cs[t])
            self.ys[t] = self.Wy @ self.hs[t] + self.by

        return self.ys

    def backward(self, targets):
        T = len(targets)

        dWf = np.zeros_like(self.Wf); dUf = np.zeros_like(self.Uf); dbf = np.zeros_like(self.bf)
        dWi = np.zeros_like(self.Wi); dUi = np.zeros_like(self.Ui); dbi = np.zeros_like(self.bi)
        dWo = np.zeros_like(self.Wo); dUo = np.zeros_like(self.Uo); dbo = np.zeros_like(self.bo)
        dWc = np.zeros_like(self.Wc); dUc = np.zeros_like(self.Uc); dbc = np.zeros_like(self.bc)
        dWy = np.zeros_like(self.Wy); dby = np.zeros_like(self.by)

        dh_next = np.zeros((self.hidden_dim, 1))
        dc_next = np.zeros((self.hidden_dim, 1))
        loss = 0.0

        for t in reversed(range(T)):
            dy = self.ys[t] - targets[t]
            loss += 0.5 * np.sum(dy ** 2)

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

            dh = self.Wy.T @ dy + dh_next

            do = dh * tanh(self.cs[t])
            do_raw = do * self.os[t] * (1 - self.os[t])

            dc = dh * self.os[t] * (1 - tanh(self.cs[t]) ** 2) + dc_next

            df = dc * self.cs[t - 1]
            df_raw = df * self.fs[t] * (1 - self.fs[t])

            di = dc * self.c_tildes[t]
            di_raw = di * self.ins[t] * (1 - self.ins[t])

            dc_tilde = dc * self.ins[t]
            dc_tilde_raw = dc_tilde * (1 - self.c_tildes[t] ** 2)

            dWf += df_raw @ self.xs[t].T; dUf += df_raw @ self.hs[t-1].T; dbf += df_raw
            dWi += di_raw @ self.xs[t].T; dUi += di_raw @ self.hs[t-1].T; dbi += di_raw
            dWo += do_raw @ self.xs[t].T; dUo += do_raw @ self.hs[t-1].T; dbo += do_raw
            dWc += dc_tilde_raw @ self.xs[t].T; dUc += dc_tilde_raw @ self.hs[t-1].T; dbc += dc_tilde_raw

            dc_next = dc * self.fs[t]
            dh_next = (self.Uf.T @ df_raw + self.Ui.T @ di_raw
                      + self.Uo.T @ do_raw + self.Uc.T @ dc_tilde_raw)

        for grad in [dWf, dUf, dbf, dWi, dUi, dbi, dWo, dUo, dbo,
                     dWc, dUc, dbc, dWy, dby]:
            np.clip(grad, -5, 5, out=grad)

        self.Wf -= self.lr * dWf; self.Uf -= self.lr * dUf; self.bf -= self.lr * dbf
        self.Wi -= self.lr * dWi; self.Ui -= self.lr * dUi; self.bi -= self.lr * dbi
        self.Wo -= self.lr * dWo; self.Uo -= self.lr * dUo; self.bo -= self.lr * dbo
        self.Wc -= self.lr * dWc; self.Uc -= self.lr * dUc; self.bc -= self.lr * dbc
        self.Wy -= self.lr * dWy; self.by -= self.lr * dby

        return loss

# ----------------------------
# データ生成: 正弦波の予測
# ----------------------------
np.random.seed(42)

# sin波にノイズを加えたデータを生成
T_total = 500
t_axis = np.linspace(0, 10 * np.pi, T_total)
data = np.sin(t_axis) + 0.1 * np.random.randn(T_total)

# 入力系列と教師系列の作成
seq_len = 20  # 過去20ステップから次の1ステップを予測
X_sequences = []
Y_sequences = []
for i in range(T_total - seq_len):
    X_sequences.append(data[i:i + seq_len])
    Y_sequences.append(data[i + seq_len])

X_sequences = np.array(X_sequences)
Y_sequences = np.array(Y_sequences)

# 訓練・テスト分割
n_train = 350
X_train = X_sequences[:n_train]
Y_train = Y_sequences[:n_train]
X_test = X_sequences[n_train:]
Y_test = Y_sequences[n_train:]

# ----------------------------
# 学習ループ
# ----------------------------
input_dim = 1
hidden_dim = 16
output_dim = 1
n_epochs = 50

gru_model = GRU(input_dim, hidden_dim, output_dim, lr=0.001)
lstm_model = LSTM(input_dim, hidden_dim, output_dim, lr=0.001)

gru_losses = []
lstm_losses = []

for epoch in range(n_epochs):
    gru_epoch_loss = 0.0
    lstm_epoch_loss = 0.0

    # ミニバッチなしで各系列ごとに学習
    for i in range(n_train):
        # 入力を (T, 1, 1) の形式に変換
        xs = [data_pt.reshape(1, 1) for data_pt in X_train[i]]
        target = [Y_train[i].reshape(1, 1)]

        # GRU
        gru_model.forward(xs)
        # 最終時刻の出力のみで損失を計算
        gru_targets = {t: np.zeros((1, 1)) for t in range(seq_len)}
        gru_targets[seq_len - 1] = target[0]
        gru_loss = gru_model.backward(gru_targets)
        gru_epoch_loss += float(np.sum((gru_model.ys[seq_len - 1] - target[0]) ** 2))

        # LSTM
        lstm_model.forward(xs)
        lstm_targets = {t: np.zeros((1, 1)) for t in range(seq_len)}
        lstm_targets[seq_len - 1] = target[0]
        lstm_loss = lstm_model.backward(lstm_targets)
        lstm_epoch_loss += float(np.sum((lstm_model.ys[seq_len - 1] - target[0]) ** 2))

    gru_losses.append(gru_epoch_loss / n_train)
    lstm_losses.append(lstm_epoch_loss / n_train)

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}/{n_epochs}  "
              f"GRU Loss: {gru_losses[-1]:.6f}  "
              f"LSTM Loss: {lstm_losses[-1]:.6f}")

# ----------------------------
# パラメータ数の確認
# ----------------------------
n_gru_params = (3 * (hidden_dim * input_dim + hidden_dim * hidden_dim + hidden_dim)
                + output_dim * hidden_dim + output_dim)
n_lstm_params = (4 * (hidden_dim * input_dim + hidden_dim * hidden_dim + hidden_dim)
                 + output_dim * hidden_dim + output_dim)
print(f"\nGRU  パラメータ数: {n_gru_params}")
print(f"LSTM パラメータ数: {n_lstm_params}")
print(f"比率 (GRU/LSTM): {n_gru_params / n_lstm_params:.2%}")

# ----------------------------
# 学習曲線の可視化
# ----------------------------
plt.figure(figsize=(10, 5))
plt.plot(range(1, n_epochs + 1), gru_losses, label="GRU", linewidth=2)
plt.plot(range(1, n_epochs + 1), lstm_losses, label="LSTM", linewidth=2, linestyle="--")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.title("Training Loss: GRU vs LSTM")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# ----------------------------
# テストデータでの予測結果の可視化
# ----------------------------
gru_preds = []
lstm_preds = []

for i in range(len(X_test)):
    xs = [data_pt.reshape(1, 1) for data_pt in X_test[i]]

    gru_model.forward(xs)
    gru_preds.append(float(gru_model.ys[seq_len - 1]))

    lstm_model.forward(xs)
    lstm_preds.append(float(lstm_model.ys[seq_len - 1]))

plt.figure(figsize=(12, 5))
plt.plot(Y_test, label="Ground Truth", linewidth=2, color="black")
plt.plot(gru_preds, label="GRU Prediction", linewidth=1.5, alpha=0.8)
plt.plot(lstm_preds, label="LSTM Prediction", linewidth=1.5, alpha=0.8, linestyle="--")
plt.xlabel("Test Sample Index")
plt.ylabel("Value")
plt.title("Time Series Prediction: GRU vs LSTM")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()