import numpy as np
import matplotlib.pyplot as plt

# ----------------------------
# ユーティリティ関数
# ----------------------------
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, axis=0, keepdims=True))
    return e / np.sum(e, axis=0, keepdims=True)

# ----------------------------
# LSTMセルの実装
# ----------------------------
class LSTMCell:
    """LSTMセル（1ステップ分の順伝播・逆伝播）"""
    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))

        # 4つのゲートのパラメータ（f, i, o, c）
        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))

    def forward(self, x, h_prev, c_prev):
        f = sigmoid(self.Wf @ x + self.Uf @ h_prev + self.bf)
        i = sigmoid(self.Wi @ x + self.Ui @ h_prev + self.bi)
        o = sigmoid(self.Wo @ x + self.Uo @ h_prev + self.bo)
        c_tilde = np.tanh(self.Wc @ x + self.Uc @ h_prev + self.bc)
        c = f * c_prev + i * c_tilde
        h = o * np.tanh(c)
        cache = (x, h_prev, c_prev, f, i, o, c_tilde, c, h)
        return h, c, cache

    def backward(self, dh, dc_next, cache):
        x, h_prev, c_prev, f, i, o, c_tilde, c, h = cache
        tanh_c = np.tanh(c)

        do = dh * tanh_c
        dc = dh * o * (1 - tanh_c ** 2) + dc_next

        df = dc * c_prev
        di = dc * c_tilde
        dc_tilde = dc * i

        df_raw = df * f * (1 - f)
        di_raw = di * i * (1 - i)
        do_raw = do * o * (1 - o)
        dc_tilde_raw = dc_tilde * (1 - c_tilde ** 2)

        # パラメータ勾配
        grads = {
            'Wf': df_raw @ x.T, 'Uf': df_raw @ h_prev.T, 'bf': df_raw,
            'Wi': di_raw @ x.T, 'Ui': di_raw @ h_prev.T, 'bi': di_raw,
            'Wo': do_raw @ x.T, 'Uo': do_raw @ h_prev.T, 'bo': do_raw,
            'Wc': dc_tilde_raw @ x.T, 'Uc': dc_tilde_raw @ h_prev.T, 'bc': dc_tilde_raw,
        }

        # 前時刻への勾配
        dh_prev = (self.Uf.T @ df_raw + self.Ui.T @ di_raw
                  + self.Uo.T @ do_raw + self.Uc.T @ dc_tilde_raw)
        dc_prev = dc * f

        return dh_prev, dc_prev, grads

# ----------------------------
# BiLSTMモデル
# ----------------------------
class BiLSTM:
    """双方向LSTM（系列ラベリング用）"""
    def __init__(self, input_dim, hidden_dim, output_dim, lr=0.01):
        self.hidden_dim = hidden_dim
        self.lr = lr

        # 順方向・逆方向のLSTMセル
        self.fwd_cell = LSTMCell(input_dim, hidden_dim)
        self.bwd_cell = LSTMCell(input_dim, hidden_dim)

        # 出力層: 2h → output_dim
        scale = np.sqrt(2.0 / (2 * hidden_dim + output_dim))
        self.Wy = np.random.randn(output_dim, 2 * hidden_dim) * scale
        self.by = np.zeros((output_dim, 1))

    def forward(self, xs):
        """順伝播: xs は (T, input_dim, 1) のリスト"""
        T = len(xs)
        h_dim = self.hidden_dim

        # 順方向
        self.fwd_caches = []
        fwd_hs = []
        h = np.zeros((h_dim, 1))
        c = np.zeros((h_dim, 1))
        for t in range(T):
            h, c, cache = self.fwd_cell.forward(xs[t], h, c)
            self.fwd_caches.append(cache)
            fwd_hs.append(h)

        # 逆方向
        self.bwd_caches = []
        bwd_hs = [None] * T
        h = np.zeros((h_dim, 1))
        c = np.zeros((h_dim, 1))
        for t in range(T - 1, -1, -1):
            h, c, cache = self.bwd_cell.forward(xs[t], h, c)
            self.bwd_caches.insert(0, cache)
            bwd_hs[t] = h

        # 結合と出力
        self.fwd_hs = fwd_hs
        self.bwd_hs = bwd_hs
        self.concat_hs = []
        self.outputs = []

        for t in range(T):
            h_concat = np.vstack([fwd_hs[t], bwd_hs[t]])  # (2h, 1)
            self.concat_hs.append(h_concat)
            logits = self.Wy @ h_concat + self.by
            probs = softmax(logits)
            self.outputs.append(probs)

        return self.outputs

    def compute_loss(self, outputs, targets):
        """クロスエントロピー損失"""
        loss = 0.0
        for t in range(len(targets)):
            loss -= np.log(outputs[t][targets[t], 0] + 1e-12)
        return loss

    def backward(self, targets):
        """逆伝播"""
        T = len(targets)
        h_dim = self.hidden_dim

        # 出力層の勾配
        dWy = np.zeros_like(self.Wy)
        dby = np.zeros_like(self.by)

        fwd_dhs = [None] * T
        bwd_dhs = [None] * T

        for t in range(T):
            dy = self.outputs[t].copy()
            dy[targets[t]] -= 1.0  # softmax + CEの勾配

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

            dh_concat = self.Wy.T @ dy  # (2h, 1)
            fwd_dhs[t] = dh_concat[:h_dim]
            bwd_dhs[t] = dh_concat[h_dim:]

        # 順方向のBPTT
        fwd_grads_accum = {}
        dh_next = np.zeros((h_dim, 1))
        dc_next = np.zeros((h_dim, 1))
        for t in range(T - 1, -1, -1):
            dh = fwd_dhs[t] + dh_next
            dh_prev, dc_prev, grads = self.fwd_cell.backward(dh, dc_next, self.fwd_caches[t])
            dh_next = dh_prev
            dc_next = dc_prev
            for key, val in grads.items():
                if key not in fwd_grads_accum:
                    fwd_grads_accum[key] = np.zeros_like(val)
                fwd_grads_accum[key] += val

        # 逆方向のBPTT（順方向に勾配を伝播）
        bwd_grads_accum = {}
        dh_next = np.zeros((h_dim, 1))
        dc_next = np.zeros((h_dim, 1))
        for t in range(T):  # t=0,1,...,T-1（逆方向RNNの逆順）
            dh = bwd_dhs[t] + dh_next
            dh_prev, dc_prev, grads = self.bwd_cell.backward(dh, dc_next, self.bwd_caches[t])
            dh_next = dh_prev
            dc_next = dc_prev
            for key, val in grads.items():
                if key not in bwd_grads_accum:
                    bwd_grads_accum[key] = np.zeros_like(val)
                bwd_grads_accum[key] += val

        # 勾配クリッピングとパラメータ更新
        for key, grad in fwd_grads_accum.items():
            np.clip(grad, -5, 5, out=grad)
            param = getattr(self.fwd_cell, key)
            param -= self.lr * grad

        for key, grad in bwd_grads_accum.items():
            np.clip(grad, -5, 5, out=grad)
            param = getattr(self.bwd_cell, key)
            param -= self.lr * grad

        np.clip(dWy, -5, 5, out=dWy)
        np.clip(dby, -5, 5, out=dby)
        self.Wy -= self.lr * dWy
        self.by -= self.lr * dby

# ----------------------------
# 単方向LSTMモデル（比較用）
# ----------------------------
class UniLSTM:
    """単方向LSTM（系列ラベリング用）"""
    def __init__(self, input_dim, hidden_dim, output_dim, lr=0.01):
        self.hidden_dim = hidden_dim
        self.lr = lr
        self.cell = LSTMCell(input_dim, hidden_dim)

        scale = np.sqrt(2.0 / (hidden_dim + output_dim))
        self.Wy = np.random.randn(output_dim, hidden_dim) * scale
        self.by = np.zeros((output_dim, 1))

    def forward(self, xs):
        T = len(xs)
        self.caches = []
        self.hs = []
        self.outputs = []

        h = np.zeros((self.hidden_dim, 1))
        c = np.zeros((self.hidden_dim, 1))
        for t in range(T):
            h, c, cache = self.cell.forward(xs[t], h, c)
            self.caches.append(cache)
            self.hs.append(h)
            logits = self.Wy @ h + self.by
            probs = softmax(logits)
            self.outputs.append(probs)

        return self.outputs

    def compute_loss(self, outputs, targets):
        loss = 0.0
        for t in range(len(targets)):
            loss -= np.log(outputs[t][targets[t], 0] + 1e-12)
        return loss

    def backward(self, targets):
        T = len(targets)
        dWy = np.zeros_like(self.Wy)
        dby = np.zeros_like(self.by)

        grads_accum = {}
        dh_next = np.zeros((self.hidden_dim, 1))
        dc_next = np.zeros((self.hidden_dim, 1))

        for t in range(T - 1, -1, -1):
            dy = self.outputs[t].copy()
            dy[targets[t]] -= 1.0

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

            dh_prev, dc_prev, grads = self.cell.backward(dh, dc_next, self.caches[t])
            dh_next = dh_prev
            dc_next = dc_prev

            for key, val in grads.items():
                if key not in grads_accum:
                    grads_accum[key] = np.zeros_like(val)
                grads_accum[key] += val

        for key, grad in grads_accum.items():
            np.clip(grad, -5, 5, out=grad)
            param = getattr(self.cell, key)
            param -= self.lr * grad

        np.clip(dWy, -5, 5, out=dWy)
        np.clip(dby, -5, 5, out=dby)
        self.Wy -= self.lr * dWy
        self.by -= self.lr * dby

# ----------------------------
# データ生成
# ----------------------------
np.random.seed(42)

# 英小文字の語彙
CHARS = list('abcdefghijklmnopqrstuvwxyz')
VOWELS = set('aeiou')
CHAR_TO_IDX = {c: i for i, c in enumerate(CHARS)}
VOCAB_SIZE = len(CHARS)
N_CLASSES = 2  # 0: 子音, 1: 母音

def generate_sequence(min_len=5, max_len=15):
    """ランダムな文字列とそのラベル（母音=1, 子音=0）を生成"""
    length = np.random.randint(min_len, max_len + 1)
    seq = [CHARS[np.random.randint(0, VOCAB_SIZE)] for _ in range(length)]
    labels = [1 if c in VOWELS else 0 for c in seq]
    return seq, labels

def seq_to_onehot(seq):
    """文字列をone-hotベクトルのリストに変換"""
    result = []
    for c in seq:
        vec = np.zeros((VOCAB_SIZE, 1))
        vec[CHAR_TO_IDX[c]] = 1.0
        result.append(vec)
    return result

# データセット生成
n_train = 500
n_test = 200

train_data = [generate_sequence() for _ in range(n_train)]
test_data = [generate_sequence() for _ in range(n_test)]

# ----------------------------
# 学習と評価
# ----------------------------
hidden_dim = 16
lr = 0.005
n_epochs = 30

bilstm = BiLSTM(VOCAB_SIZE, hidden_dim, N_CLASSES, lr=lr)
unilstm = UniLSTM(VOCAB_SIZE, hidden_dim, N_CLASSES, lr=lr)

bi_losses = []
uni_losses = []
bi_accs = []
uni_accs = []

for epoch in range(n_epochs):
    bi_epoch_loss = 0.0
    uni_epoch_loss = 0.0

    # シャッフル
    indices = np.random.permutation(n_train)

    for idx in indices:
        seq, labels = train_data[idx]
        xs = seq_to_onehot(seq)

        # BiLSTM
        bi_out = bilstm.forward(xs)
        bi_loss = bilstm.compute_loss(bi_out, labels)
        bilstm.backward(labels)
        bi_epoch_loss += bi_loss

        # 単方向LSTM
        uni_out = unilstm.forward(xs)
        uni_loss = unilstm.compute_loss(uni_out, labels)
        unilstm.backward(labels)
        uni_epoch_loss += uni_loss

    bi_losses.append(bi_epoch_loss / n_train)
    uni_losses.append(uni_epoch_loss / n_train)

    # テスト精度の計算
    bi_correct = 0
    uni_correct = 0
    total = 0

    for seq, labels in test_data:
        xs = seq_to_onehot(seq)

        bi_out = bilstm.forward(xs)
        uni_out = unilstm.forward(xs)

        for t in range(len(labels)):
            bi_pred = int(np.argmax(bi_out[t]))
            uni_pred = int(np.argmax(uni_out[t]))
            if bi_pred == labels[t]:
                bi_correct += 1
            if uni_pred == labels[t]:
                uni_correct += 1
            total += 1

    bi_accs.append(bi_correct / total)
    uni_accs.append(uni_correct / total)

    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1}/{n_epochs}  "
              f"BiLSTM Loss: {bi_losses[-1]:.4f} Acc: {bi_accs[-1]:.4f}  "
              f"UniLSTM Loss: {uni_losses[-1]:.4f} Acc: {uni_accs[-1]:.4f}")

# ----------------------------
# 結果の可視化
# ----------------------------
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

# 損失曲線
ax1.plot(range(1, n_epochs + 1), bi_losses, label="BiLSTM", linewidth=2)
ax1.plot(range(1, n_epochs + 1), uni_losses, label="UniLSTM", linewidth=2, linestyle="--")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Cross-Entropy Loss")
ax1.set_title("Training Loss: BiLSTM vs UniLSTM")
ax1.legend()
ax1.grid(True, alpha=0.3)

# 精度曲線
ax2.plot(range(1, n_epochs + 1), bi_accs, label="BiLSTM", linewidth=2)
ax2.plot(range(1, n_epochs + 1), uni_accs, label="UniLSTM", linewidth=2, linestyle="--")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Accuracy")
ax2.set_title("Test Accuracy: BiLSTM vs UniLSTM")
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# ----------------------------
# 推論例の表示
# ----------------------------
print("\n--- 推論例（BiLSTM）---")
for i in range(5):
    seq, labels = test_data[i]
    xs = seq_to_onehot(seq)
    outputs = bilstm.forward(xs)

    preds = [int(np.argmax(outputs[t])) for t in range(len(seq))]
    print(f"入力: {''.join(seq)}")
    print(f"正解: {labels}")
    print(f"予測: {preds}")
    print(f"一致: {['o' if p == l else 'x' for p, l in zip(preds, labels)]}")
    print()