import csv
import numpy as np
from tqdm import tqdm

# --- 유틸 함수들 ---
def gelu(x):
    return 0.5 * x * (1 + np.tanh(np.sqrt(2/np.pi) * (x + 0.044715 * x**3)))

def swiglu(x):
    x1, x2 = np.split(x, 2, axis=-1)
    return x1 * gelu(x2)

def softmax(x):
    x = x - np.max(x, axis=-1, keepdims=True)
    e_x = np.exp(x)
    return e_x / e_x.sum(axis=-1, keepdims=True)

def top_p_sample(probs, p=0.9):
    sorted_indices = np.argsort(probs)[::-1]
    sorted_probs = probs[sorted_indices]
    cumulative = np.cumsum(sorted_probs)
    cutoff = np.searchsorted(cumulative, p) + 1
    top_indices = sorted_indices[:cutoff]
    top_probs = probs[top_indices]
    top_probs /= top_probs.sum()
    return np.random.choice(top_indices, p=top_probs)

def scaled_dot_product_attention(Q, K, V, causal_mask=True):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)  # (batch, seq_len_q, seq_len_k)

    if causal_mask:
        seq_len_q, seq_len_k = scores.shape[1], scores.shape[2]
        # 아래 삼각행렬 형태로 마스크 생성 (True는 마스킹할 위치)
        mask = np.triu(np.ones((seq_len_q, seq_len_k), dtype=bool), k=1)
        # mask를 (1, seq_len_q, seq_len_k)로 broadcasting 하기 위해 차원 추가
        mask = mask[None, :, :]  # (1, seq_len, seq_len)
        
        scores = np.where(mask, -1e9, scores)  # 마스크 위치에 매우 큰 음수 넣기

    scores -= np.max(scores, axis=-1, keepdims=True)
    weights = np.exp(scores)
    weights /= weights.sum(axis=-1, keepdims=True)
    output = weights @ V
    return output, weights


def scaled_dot_product_attention_backward(Q, K, V, weights, grad_output, causal_mask=True):
    d_k = Q.shape[-1]
    grad_weights = grad_output @ V.transpose(0, 2, 1)
    grad_V = weights.transpose(0, 2, 1) @ grad_output

    s = weights
    grad_scores = s * (grad_weights - (grad_weights * s).sum(axis=-1, keepdims=True))

    if causal_mask:
        seq_len_q, seq_len_k = grad_scores.shape[1], grad_scores.shape[2]
        mask = np.triu(np.ones((seq_len_q, seq_len_k), dtype=bool), k=1)
        mask = mask[None, :, :]
        # 마스크 위치는 grad_scores 0으로 만들어서 영향 안 미치게 처리
        grad_scores = np.where(mask, 0, grad_scores)

    grad_Q = grad_scores @ K / np.sqrt(d_k)
    grad_K = grad_scores.transpose(0, 2, 1) @ Q / np.sqrt(d_k)

    return grad_Q, grad_K, grad_V


# --- Adam Optimizer ---
class AdamOptimizer:
    def __init__(self, params, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8):
        self.params = params
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = [np.zeros_like(p) for p in params]
        self.v = [np.zeros_like(p) for p in params]
        self.t = 0

    def step(self, grads):
        self.t += 1
        lr_t = self.lr * (np.sqrt(1 - self.beta2 ** self.t) / (1 - self.beta1 ** self.t))
        for i in range(len(self.params)):
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grads[i]
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * (grads[i] ** 2)
            m_hat = self.m[i] / (1 - self.beta1 ** self.t)
            v_hat = self.v[i] / (1 - self.beta2 ** self.t)
            self.params[i] -= lr_t * m_hat / (np.sqrt(v_hat) + self.eps)

# --- LayerNorm 클래스 ---
class LayerNorm:
    def __init__(self, dim, eps=1e-5):
        self.eps = eps
        self.gamma = np.ones(dim)
        self.beta = np.zeros(dim)

    def forward(self, x):
        self.mean = x.mean(axis=-1, keepdims=True)
        self.var = x.var(axis=-1, keepdims=True)
        self.x_centered = x - self.mean
        self.std = np.sqrt(self.var + self.eps)
        self.x_norm = self.x_centered / self.std
        return self.gamma * self.x_norm + self.beta

    def backward(self, grad_out, lr):
        N = grad_out.shape[-1]
        grad_gamma = (grad_out * self.x_norm).sum(axis=(0, 1))
        grad_beta = grad_out.sum(axis=(0, 1))

        dx_norm = grad_out * self.gamma
        dvar = (-0.5) * (dx_norm * self.x_centered).sum(axis=-1, keepdims=True) * (self.var + self.eps) ** (-1.5)
        dmean = (-dx_norm / self.std).sum(axis=-1, keepdims=True) + dvar * (-2.0 / N) * self.x_centered.sum(axis=-1, keepdims=True)
        dx = dx_norm / self.std + dvar * 2 * self.x_centered / N + dmean / N

        self.gamma -= lr * grad_gamma
        self.beta -= lr * grad_beta
        return dx

# --- EmbeddingLayer ---
class EmbeddingLayer:
    def __init__(self, vocab_size, embed_dim, seed=0):
        rng = np.random.default_rng(seed)
        self.W = rng.normal(0, 0.1, (vocab_size, embed_dim))

    def forward(self, indices):
        self.indices = indices
        return self.W[indices]

    def backward(self, grad_out, lr=0.01):
        flat_indices = self.indices.flatten()
        grad_flat = grad_out.reshape(-1, self.W.shape[1])
        np.add.at(self.W, flat_indices, -lr * grad_flat)

# --- Xavier 초기화 ---
def xavier_init(fan_in, fan_out):
    limit = np.sqrt(6.0 / (fan_in + fan_out))
    return np.random.uniform(-limit, limit, size=(fan_in, fan_out))

class GLULayer:
    def __init__(self, input_dim, hidden_dim, output_dim, seed=0):
        rng = np.random.default_rng(seed)

        # GLU 레이어 파라미터
        self.W1 = xavier_init(hidden_dim, hidden_dim * 2)
        self.b1 = np.zeros(hidden_dim * 2)


        self.W2 = xavier_init(hidden_dim, output_dim)     # hidden_dim → output_dim
        self.b2 = np.zeros(output_dim)

        # Attention용 프로젝션 행렬들
        self.Wq_proj = xavier_init(input_dim, hidden_dim)  # 입력을 hidden_dim으로 변환 (attention 입력)
        self.bq_proj = np.zeros(hidden_dim)

        self.Wq = xavier_init(hidden_dim, hidden_dim)      # Q 생성용
        self.Wk = xavier_init(input_dim, hidden_dim)       # K 생성용
        self.Wv = xavier_init(input_dim, hidden_dim)       # V 생성용

        self.layernorm = LayerNorm(hidden_dim)

        # Adam 모멘텀 변수 초기화
        self.m_W1 = np.zeros_like(self.W1)
        self.v_W1 = np.zeros_like(self.W1)
        self.m_b1 = np.zeros_like(self.b1)
        self.v_b1 = np.zeros_like(self.b1)

        self.m_W2 = np.zeros_like(self.W2)
        self.v_W2 = np.zeros_like(self.W2)
        self.m_b2 = np.zeros_like(self.b2)
        self.v_b2 = np.zeros_like(self.b2)

        self.m_Wq_proj = np.zeros_like(self.Wq_proj)
        self.v_Wq_proj = np.zeros_like(self.Wq_proj)
        self.m_bq_proj = np.zeros_like(self.bq_proj)
        self.v_bq_proj = np.zeros_like(self.bq_proj)

        self.m_Wq = np.zeros_like(self.Wq)
        self.v_Wq = np.zeros_like(self.Wq)
        self.m_Wk = np.zeros_like(self.Wk)
        self.v_Wk = np.zeros_like(self.Wk)
        self.m_Wv = np.zeros_like(self.Wv)
        self.v_Wv = np.zeros_like(self.Wv)

        self.beta1 = 0.9
        self.beta2 = 0.999
        self.eps = 1e-8
        self.t = 0

    def attention(self, h_proj, x):
        Q = h_proj @ self.Wq     # (batch, seq_len, hidden_dim) @ (hidden_dim, hidden_dim)
        K = x @ self.Wk          # (batch, seq_len, input_dim) @ (input_dim, hidden_dim)
        V = x @ self.Wv          # (batch, seq_len, input_dim) @ (input_dim, hidden_dim)
        attn_out, weights = scaled_dot_product_attention(Q, K, V)
        self.cache = {'Q': Q, 'K': K, 'V': V, 'weights': weights, 'h_proj': h_proj, 'x': x}
        return attn_out


    def forward(self, x):
        self.x = x  # input shape: (batch, seq_len, input_dim)

        # input_dim → hidden_dim로 변환 (attention 입력용)
        self.h_proj = x @ self.Wq_proj + self.bq_proj  # (batch, seq_len, hidden_dim)

        # attention 수행
        self.attn_out = self.attention(self.h_proj, self.x)
 # output shape: (batch, seq_len, hidden_dim)

        # GLU 전처리
        self.h_in = self.attn_out @ self.W1 + self.b1  # (batch, seq_len, 2*hidden_dim)
        x1, x2 = np.split(self.h_in, 2, axis=-1)
        self.relu_mask = (x2 > 0).astype(float)  # ReLU mask 저장 (swiglu에 사용)

        # swiglu 활성화 + LayerNorm
        self.h = swiglu(self.h_in)
        self.h = self.layernorm.forward(self.h)

        # residual 연결
        self.combined = self.attn_out + self.h       # (batch, seq_len, hidden_dim)

        self.out_seq = self.combined @ self.W2 + self.b2  # (batch, seq_len, output_dim)
        self.out = self.out_seq[:, -1, :]  # 마지막 시퀀스 출력 (batch, output_dim)

        return self.out

    def predict_proba(self, x):
        logits = self.forward(x)
        return softmax(logits)

    def backward(self, grad_out, lr=0.05):
        self.t += 1
        batch_size = grad_out.shape[0]

    # W2, b2 gradient (마지막 시퀀스에 대해서만)
        grad_W2 = (self.combined[:, -1, :].T @ grad_out) / batch_size
        grad_b2 = grad_out.mean(axis=0)

    # grad_combined: (batch, seq_len, hidden_dim), 마지막 타임스텝에만 grad_out 영향
        grad_combined_last = grad_out @ self.W2.T
        grad_combined = np.zeros_like(self.combined)
        grad_combined[:, -1, :] = grad_combined_last

    # residual 연결: combined = attn_out + h
        grad_attn_out = grad_combined.copy()
        grad_h = grad_combined.copy()

    # Attention backward 준비
        Q = self.cache['Q']
        K = self.cache['K']
        V = self.cache['V']
        weights = self.cache['weights']
        h_proj = self.cache['h_proj']
        x = self.cache['x']

        grad_Q, grad_K, grad_V = scaled_dot_product_attention_backward(Q, K, V, weights, grad_attn_out)

    # Wq, Wk, Wv gradient 계산
        grad_Wq = (h_proj.transpose(0, 2, 1) @ grad_Q).sum(axis=0) / batch_size
        grad_Wk = (x.transpose(0, 2, 1) @ grad_K).sum(axis=0) / batch_size
        grad_Wv = (x.transpose(0, 2, 1) @ grad_V).sum(axis=0) / batch_size

    # h_proj에 대한 gradient만 grad_h에 더하기
        grad_h += grad_Q @ self.Wq.T

    # x에 대한 gradient는 별도로 계산
        grad_x_from_K = grad_K @ self.Wk.T
        grad_x_from_V = grad_V @ self.Wv.T

    # LayerNorm backward
        grad_h = self.layernorm.backward(grad_h, lr)

    # swiglu backward
        x1, x2 = np.split(self.h_in, 2, axis=-1)
        relu_mask = self.relu_mask

        grad_x1 = grad_h * np.maximum(0, x2)
        grad_x2 = grad_h * x1 * relu_mask
        grad_h_in = np.concatenate([grad_x1, grad_x2], axis=-1)

    # W1, b1 gradient
        grad_W1 = (self.attn_out.transpose(0, 2, 1) @ grad_h_in).sum(axis=0) / batch_size
        grad_b1 = grad_h_in.mean(axis=(0, 1))

    # attn_out에 대한 gradient (GLU 통과 후)
        grad_attn_out_from_glu = grad_h_in @ self.W1.T

    # 전체 x gradient = attn_out gradient + K,V로부터 온 gradient
        grad_x = grad_attn_out + grad_attn_out_from_glu + grad_x_from_K + grad_x_from_V

    # Wq_proj, bq_proj gradient 계산
        grad_Wq_proj = (x.transpose(0, 2, 1) @ grad_x).sum(axis=0) / batch_size
        grad_bq_proj = grad_x.mean(axis=(0, 1))

    # input x에 대한 gradient
        grad_input = grad_x @ self.Wq_proj.T

    # Adam 업데이트 함수 (내부 정의)
        def adam_update(param, grad, m, v, t, lr, beta1=0.9, beta2=0.999, eps=1e-8):
            m = beta1 * m + (1 - beta1) * grad
            v = beta2 * v + (1 - beta2) * (grad ** 2)
            m_hat = m / (1 - beta1 ** t)
            v_hat = v / (1 - beta2 ** t)
            param_update = lr * m_hat / (np.sqrt(v_hat) + eps)
            param -= param_update
            return param, m, v

    # 파라미터 업데이트
        self.W2, self.m_W2, self.v_W2 = adam_update(self.W2, grad_W2, self.m_W2, self.v_W2, self.t, lr)
        self.b2, self.m_b2, self.v_b2 = adam_update(self.b2, grad_b2, self.m_b2, self.v_b2, self.t, lr)

        self.Wq, self.m_Wq, self.v_Wq = adam_update(self.Wq, grad_Wq, self.m_Wq, self.v_Wq, self.t, lr)
        self.Wk, self.m_Wk, self.v_Wk = adam_update(self.Wk, grad_Wk, self.m_Wk, self.v_Wk, self.t, lr)
        self.Wv, self.m_Wv, self.v_Wv = adam_update(self.Wv, grad_Wv, self.m_Wv, self.v_Wv, self.t, lr)

        self.W1, self.m_W1, self.v_W1 = adam_update(self.W1, grad_W1, self.m_W1, self.v_W1, self.t, lr)
        self.b1, self.m_b1, self.v_b1 = adam_update(self.b1, grad_b1, self.m_b1, self.v_b1, self.t, lr)

        self.Wq_proj, self.m_Wq_proj, self.v_Wq_proj = adam_update(self.Wq_proj, grad_Wq_proj, self.m_Wq_proj, self.v_Wq_proj, self.t, lr)
        self.bq_proj, self.m_bq_proj, self.v_bq_proj = adam_update(self.bq_proj, grad_bq_proj, self.m_bq_proj, self.v_bq_proj, self.t, lr)

        return grad_input


# --- 손실 함수 ---
def sparse_cross_entropy_loss(probs, true_indices):
    batch = probs.shape[0]
    clipped_probs = np.clip(probs, 1e-9, 1.0)
    log_probs = np.log(clipped_probs)
    loss = -np.sum(log_probs[np.arange(batch), true_indices]) / batch
    return loss

def grad_sparse_cross_entropy(probs, true_indices):
    grad = probs.copy()
    grad[np.arange(probs.shape[0]), true_indices] -= 1
    grad /= probs.shape[0]
    return grad


def build_vocab(pairs):
    tokens = set()
    for inp, out in pairs:
        tokens.update(inp.split())
        tokens.update(out.split())
    tokens.add('<EOS>')
    token2idx = {tok: idx for idx, tok in enumerate(sorted(tokens))}
    idx2token = {idx: tok for tok, idx in token2idx.items()}
    return token2idx, idx2token

def tokenize_and_encode(pairs, token2idx):
    contexts = []
    next_tokens = []
    for inp_text, out_text in pairs:
        inp_tokens = inp_text.split()
        out_tokens = out_text.split() + ['<EOS>']
        seq = inp_tokens.copy()
        for token in out_tokens:
            contexts.append(seq.copy())
            next_tokens.append(token)
            seq.append(token)
    contexts_idx = [[token2idx[tok] for tok in c] for c in contexts]
    next_tokens_idx = [token2idx[tok] for tok in next_tokens]
    max_len = max(len(c) for c in contexts_idx)
    X_padded = np.array([np.pad(c, (0, max_len - len(c)), constant_values=0) for c in contexts_idx])
    Y_indices = np.array(next_tokens_idx)
    return X_padded, Y_indices

def generate_sequence(model, token2idx, idx2token, prompt, max_tokens=100, top_p=0.75, min_length=5, temperature=0.7):
    if isinstance(prompt, str):
        prompt_tokens = prompt.split()
    else:
        prompt_tokens = prompt

    generated = []
    token_counts = {}

    for step in range(max_tokens):
        context = prompt_tokens + generated
        try:
            X_indices = np.array([[token2idx[tok] for tok in context]])
        except KeyError:
            print(f"[Warning] Unknown token in prompt/context: {context}")
            break

        probs = model.predict(X_indices)
        probs = np.clip(probs, 1e-9, None)

        # 온도 적용
        probs = np.power(probs, 1.0 / temperature)
        probs /= probs.sum()

        # 반복 토큰 억제 (count에 비례해서 페널티 더 강하게)
        for token, count in token_counts.items():
            idx = token2idx[token]
            probs[idx] /= (count + 1) ** 2

        probs /= probs.sum()

        next_idx = top_p_sample(probs, p=top_p)
        next_token = idx2token[next_idx]

        if next_token == '<EOS>' and step < min_length:
            probs[next_idx] = 0
            probs /= probs.sum()
            next_idx = top_p_sample(probs, p=top_p)
            next_token = idx2token[next_idx]

        if next_token == '<EOS>':
            break

        generated.append(next_token)
        token_counts[next_token] = token_counts.get(next_token, 0) + 1

    return generated


class Model:
    def __init__(self, n_models, vocab_size, embed_dim, hidden_dim, output_dim):
        self.embedding = EmbeddingLayer(vocab_size, embed_dim)
        self.models = [GLULayer(embed_dim, hidden_dim, output_dim, seed=i) for i in range(n_models)]
        # 이 두 개를 멤버 변수로 저장해서 나중에 쓸 수 있게!
        self.token2idx = None
        self.idx2token = None

    def fit(self, X_indices, Y_indices, token2idx, idx2token, epochs=30, lr=0.003, batch_size=32):
        self.token2idx = token2idx
        self.idx2token = idx2token

        n_samples = X_indices.shape[0]
        n_batches = (n_samples + batch_size - 1) // batch_size

        for i, model in enumerate(self.models):
            print(f"Training model {i+1}/{len(self.models)}")
            for epoch in range(epochs):
                epoch_loss = 0.0

                with tqdm(range(n_batches), desc=f"Epoch {epoch+1}/{epochs}", unit="batch") as pbar:
                    for batch_idx in pbar:
                        start = batch_idx * batch_size
                        end = min(start + batch_size, n_samples)

                        batch_X = X_indices[start:end]
                        batch_Y = Y_indices[start:end]

                        embeddings = self.embedding.forward(batch_X)
                        logits = model.forward(embeddings)
                        probs = softmax(logits)

                        loss = sparse_cross_entropy_loss(probs, batch_Y)
                        epoch_loss += loss * (end - start)

                        grad_out = grad_sparse_cross_entropy(probs, batch_Y)
                        grad_embedding = model.backward(grad_out, lr)
                        self.embedding.backward(grad_embedding, lr)

                        pbar.set_postfix(loss=f"{loss:.4f}")

                avg_epoch_loss = epoch_loss / n_samples
                print(f"Epoch {epoch+1} average loss: {avg_epoch_loss:.4f}")

# fit 함수 내 epoch 끝난 뒤
                sample_tokens = generate_sequence(self, self.token2idx, self.idx2token, "건강한")
                sample_sentence = " ".join(sample_tokens)
                print(f"[Epoch {epoch+1}] Sample: {sample_sentence}")

    def predict(self, X_indices):
        embeddings = self.embedding.forward(X_indices)
        preds = np.array([model.predict_proba(embeddings) for model in self.models])
        idx = np.random.choice(len(self.models))
        return preds[idx][0]