import numpy as np
import matplotlib.pyplot as plt

# ----------------------------
# 定数と特殊トークン
# ----------------------------
PAD = 0
SOS = 10
EOS = 11
VOCAB_SIZE = 12  # 0-9, SOS, EOS

# ----------------------------
# データ生成: 数列逆順タスク
# ----------------------------
def generate_data(n_samples, min_len=3, max_len=6):
    """ランダムな数列とその逆順をペアで生成"""
    data = []
    for _ in range(n_samples):
        length = np.random.randint(min_len, max_len + 1)
        seq = np.random.randint(1, 10, size=length).tolist()  # 1-9の数列
        target = seq[::-1]  # 逆順
        data.append((seq, target))
    return data

# ----------------------------
# ユーティリティ関数
# ----------------------------
def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

def softmax(x):
    e = np.exp(x - np.max(x))
    return e / np.sum(e)

def one_hot(idx, size):
    """ワンホットベクトルを生成"""
    vec = np.zeros((size, 1))
    vec[idx] = 1.0
    return vec

# ----------------------------
# GRUセル（Encoder/Decoder共通）
# ----------------------------
class GRUCell:
    def __init__(self, input_dim, hidden_dim):
        self.hidden_dim = hidden_dim
        scale_ih = np.sqrt(2.0 / (input_dim + hidden_dim))
        scale_hh = np.sqrt(2.0 / (hidden_dim + hidden_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))

    def forward(self, x, h_prev):
        """1ステップの順伝播"""
        z = sigmoid(self.Wz @ x + self.Uz @ h_prev + self.bz)
        r = sigmoid(self.Wr @ x + self.Ur @ h_prev + self.br)
        h_tilde = np.tanh(self.Wh @ x + self.Uh @ (r * h_prev) + self.bh)
        h = z * h_prev + (1 - z) * h_tilde
        return h, z, r, h_tilde

    def get_params(self):
        return [self.Wz, self.Uz, self.bz, self.Wr, self.Ur, self.br,
                self.Wh, self.Uh, self.bh]

# ----------------------------
# Seq2Seqモデル
# ----------------------------
class Seq2Seq:
    def __init__(self, vocab_size, hidden_dim, lr=0.01):
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.lr = lr

        # Encoder GRUセル
        self.encoder = GRUCell(vocab_size, hidden_dim)

        # Decoder GRUセル
        self.decoder = GRUCell(vocab_size, hidden_dim)

        # 出力層: 隠れ状態 → 語彙分布
        scale = np.sqrt(2.0 / (hidden_dim + vocab_size))
        self.W_out = np.random.randn(vocab_size, hidden_dim) * scale
        self.b_out = np.zeros((vocab_size, 1))

    def encode(self, input_seq):
        """Encoder: 入力系列 → コンテキストベクトル"""
        h = np.zeros((self.hidden_dim, 1))
        self.enc_states = []

        for token in input_seq:
            x = one_hot(token, self.vocab_size)
            h, z, r, h_tilde = self.encoder.forward(x, h)
            self.enc_states.append({
                'x': x, 'h': h, 'z': z, 'r': r, 'h_tilde': h_tilde,
                'h_prev': self.enc_states[-1]['h'] if self.enc_states else np.zeros_like(h)
            })

        return h  # コンテキストベクトル

    def decode_train(self, context, target_seq):
        """Decoder（Teacher Forcing使用）: 学習時"""
        h = context
        self.dec_states = []
        loss = 0.0

        # SOS + target_seq がDecoderへの入力
        input_tokens = [SOS] + target_seq

        for t in range(len(target_seq)):
            x = one_hot(input_tokens[t], self.vocab_size)
            h_prev = h
            h, z, r, h_tilde = self.decoder.forward(x, h)

            # 出力確率の計算
            logits = self.W_out @ h + self.b_out
            probs = softmax(logits)

            # クロスエントロピー損失
            target_token = target_seq[t]
            loss -= np.log(probs[target_token, 0] + 1e-12)

            self.dec_states.append({
                'x': x, 'h': h, 'h_prev': h_prev, 'z': z, 'r': r,
                'h_tilde': h_tilde, 'logits': logits, 'probs': probs,
                'target': target_token
            })

        return loss

    def decode_inference(self, context, max_len=10):
        """Decoder（推論時）: 貪欲デコーディング"""
        h = context
        output = []
        token = SOS

        for _ in range(max_len):
            x = one_hot(token, self.vocab_size)
            h, _, _, _ = self.decoder.forward(x, h)
            logits = self.W_out @ h + self.b_out
            probs = softmax(logits)
            token = int(np.argmax(probs))

            if token == EOS:
                break
            output.append(token)

        return output

    def backward(self, input_seq, target_seq):
        """逆伝播: 数値勾配で簡易実装"""
        eps = 1e-5
        all_params = (self.encoder.get_params()
                     + self.decoder.get_params()
                     + [self.W_out, self.b_out])
        grads = []

        for param in all_params:
            grad = np.zeros_like(param)
            it = np.nditer(param, flags=['multi_index'])
            while not it.finished:
                idx = it.multi_index

                # f(x + eps)
                param[idx] += eps
                context_plus = self.encode(input_seq)
                loss_plus = self.decode_train(context_plus, target_seq)

                # f(x - eps)
                param[idx] -= 2 * eps
                context_minus = self.encode(input_seq)
                loss_minus = self.decode_train(context_minus, target_seq)

                # 中心差分
                grad[idx] = (loss_plus - loss_minus) / (2 * eps)
                param[idx] += eps  # 元に戻す

                it.iternext()
            grads.append(grad)

        # パラメータ更新
        for param, grad in zip(all_params, grads):
            np.clip(grad, -5, 5, out=grad)
            param -= self.lr * grad

    def train_analytic(self, input_seq, target_seq):
        """解析的な逆伝播による学習（出力層のみ高速化）"""
        # 順伝播
        context = self.encode(input_seq)
        loss = self.decode_train(context, target_seq)

        # Decoder出力層の解析的勾配
        dW_out = np.zeros_like(self.W_out)
        db_out = np.zeros_like(self.b_out)

        # Decoder隠れ状態の勾配
        dh_next = np.zeros((self.hidden_dim, 1))

        # 各Decoderタイムステップを逆順に辿る
        dec_dhs = []
        for t in reversed(range(len(self.dec_states))):
            state = self.dec_states[t]

            # softmaxの勾配
            dy = state['probs'].copy()
            dy[state['target']] -= 1.0  # ∂L/∂logits

            dW_out += dy @ state['h'].T
            db_out += dy
            dh = self.W_out.T @ dy + dh_next

            # GRUの逆伝播
            z = state['z']
            r = state['r']
            h_tilde = state['h_tilde']
            h_prev = state['h_prev']

            dz = dh * (h_prev - h_tilde)
            dz_raw = dz * z * (1 - z)

            dh_tilde = dh * (1 - z)
            dh_tilde_raw = dh_tilde * (1 - h_tilde ** 2)

            dr = (self.decoder.Uh.T @ dh_tilde_raw) * h_prev
            dr_raw = dr * r * (1 - r)

            # Decoder GRUのパラメータ勾配
            self.decoder.Wz -= self.lr * (dz_raw @ state['x'].T)
            self.decoder.Uz -= self.lr * (dz_raw @ h_prev.T)
            self.decoder.bz -= self.lr * dz_raw
            self.decoder.Wr -= self.lr * (dr_raw @ state['x'].T)
            self.decoder.Ur -= self.lr * (dr_raw @ h_prev.T)
            self.decoder.br -= self.lr * dr_raw
            self.decoder.Wh -= self.lr * (dh_tilde_raw @ state['x'].T)
            self.decoder.Uh -= self.lr * (dh_tilde_raw @ (r * h_prev).T)
            self.decoder.bh -= self.lr * dh_tilde_raw

            # 前時刻への勾配伝播
            dh_next = (dh * z
                      + self.decoder.Uz.T @ dz_raw
                      + self.decoder.Ur.T @ dr_raw
                      + (self.decoder.Uh.T @ dh_tilde_raw) * r)

            dec_dhs.append(dh_next)

        # 出力層パラメータ更新
        np.clip(dW_out, -5, 5, out=dW_out)
        np.clip(db_out, -5, 5, out=db_out)
        self.W_out -= self.lr * dW_out
        self.b_out -= self.lr * db_out

        # Encoderの逆伝播（dh_nextをEncoderに伝播）
        dh_enc = dh_next
        for t in reversed(range(len(self.enc_states))):
            state = self.enc_states[t]
            z = state['z']
            r = state['r']
            h_tilde = state['h_tilde']
            h_prev = state['h_prev']

            dz = dh_enc * (h_prev - h_tilde)
            dz_raw = dz * z * (1 - z)

            dh_tilde = dh_enc * (1 - z)
            dh_tilde_raw = dh_tilde * (1 - h_tilde ** 2)

            dr = (self.encoder.Uh.T @ dh_tilde_raw) * h_prev
            dr_raw = dr * r * (1 - r)

            self.encoder.Wz -= self.lr * (dz_raw @ state['x'].T)
            self.encoder.Uz -= self.lr * (dz_raw @ h_prev.T)
            self.encoder.bz -= self.lr * dz_raw
            self.encoder.Wr -= self.lr * (dr_raw @ state['x'].T)
            self.encoder.Ur -= self.lr * (dr_raw @ h_prev.T)
            self.encoder.br -= self.lr * dr_raw
            self.encoder.Wh -= self.lr * (dh_tilde_raw @ state['x'].T)
            self.encoder.Uh -= self.lr * (dh_tilde_raw @ (r * h_prev).T)
            self.encoder.bh -= self.lr * dh_tilde_raw

            dh_enc = (dh_enc * z
                     + self.encoder.Uz.T @ dz_raw
                     + self.encoder.Ur.T @ dr_raw
                     + (self.encoder.Uh.T @ dh_tilde_raw) * r)

        return loss

# ----------------------------
# 学習の実行
# ----------------------------
np.random.seed(42)

# ハイパーパラメータ
hidden_dim = 32
lr = 0.005
n_epochs = 100
n_train = 500
n_test = 100

# データ生成
train_data = generate_data(n_train, min_len=3, max_len=5)
test_data = generate_data(n_test, min_len=3, max_len=5)

# モデル初期化
model = Seq2Seq(VOCAB_SIZE, hidden_dim, lr=lr)

# 学習ループ
losses = []
accuracies = []

for epoch in range(n_epochs):
    epoch_loss = 0.0
    np.random.shuffle(train_data)

    for input_seq, target_seq in train_data:
        target_with_eos = target_seq + [EOS]
        loss = model.train_analytic(input_seq, target_with_eos)
        epoch_loss += loss

    avg_loss = epoch_loss / n_train
    losses.append(avg_loss)

    # テスト精度の計算
    if (epoch + 1) % 10 == 0:
        correct = 0
        for input_seq, target_seq in test_data:
            context = model.encode(input_seq)
            pred = model.decode_inference(context, max_len=len(target_seq) + 2)
            if pred == target_seq:
                correct += 1
        acc = correct / n_test
        accuracies.append((epoch + 1, acc))
        print(f"Epoch {epoch+1}/{n_epochs}  Loss: {avg_loss:.4f}  "
              f"Test Accuracy: {acc:.2%}")

# ----------------------------
# 学習曲線の可視化
# ----------------------------
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

# 損失曲線
ax1.plot(range(1, n_epochs + 1), losses, linewidth=2)
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Cross-Entropy Loss")
ax1.set_title("Training Loss")
ax1.grid(True, alpha=0.3)

# 精度曲線
if accuracies:
    epochs_acc, accs = zip(*accuracies)
    ax2.plot(epochs_acc, accs, 'o-', linewidth=2, markersize=8)
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Sequence Accuracy")
    ax2.set_title("Test Accuracy (Exact Match)")
    ax2.set_ylim([0, 1.05])
    ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# ----------------------------
# 推論例の表示
# ----------------------------
print("\n--- 推論例 ---")
for i in range(10):
    input_seq, target_seq = test_data[i]
    context = model.encode(input_seq)
    pred = model.decode_inference(context, max_len=len(target_seq) + 2)
    status = "OK" if pred == target_seq else "NG"
    print(f"入力: {input_seq}  正解: {target_seq}  予測: {pred}  [{status}]")