# Flax + JAX TPU-ready reimplementation of your ReLM model and training loop.
# Requirements:
# pip install --upgrade "jax[tpu]" flax optax sentencepiece

import os
import math
import numpy as np
import sentencepiece as spm
from functools import partial
from typing import Any, Callable, Optional, Tuple, Sequence
import requests
import jax
import jax.numpy as jnp
from jax import random
from flax import linen as nn
from flax.training import train_state, checkpoints
import optax
import tqdm

def download_file(url, save_path):
    r = requests.get(url, stream=True)
    r.raise_for_status()
    with open(save_path, "wb") as f:
        for chunk in r.iter_content(8192*2):
            f.write(chunk)
    print(f"✅ {save_path} 저장됨")
# ------------------
# Config
# ------------------
SEQ_LEN = 512
# global batch size (across all devices)
GLOBAL_BATCH = 256
# adjust for memory
LIMIT = 200_000         # number of sequences to load (reduce if OOM)
VOCAB_MODEL = "ko_unigram.model"
CORPUS_PATH = "corpus.txt"
DTYPE = jnp.bfloat16 if jax.local_devices()[0].platform == "tpu" else jnp.float32
SEED = 42
LEARNING_RATE = 1e-4
EPOCHS = 1

if not os.path.exists(CORPUS_PATH):
    download_file(
        "https://huggingface.co/datasets/Yuchan5386/Prototype/resolve/main/corpus_ko.txt?download=true",
        CORPUS_PATH
    )

if not os.path.exists(VOCAB_MODEL):
    download_file(
        "https://huggingface.co/Yuchan5386/inlam-100m/resolve/main/ko_unigram.model?download=true",
        VOCAB_MODEL
    )

# Derived
NUM_DEVICES = jax.device_count()
assert GLOBAL_BATCH % NUM_DEVICES == 0, "GLOBAL_BATCH must be divisible by device count"
PER_DEVICE_BATCH = GLOBAL_BATCH // NUM_DEVICES

print("devices:", jax.devices())
print("num_devices:", NUM_DEVICES, "per_device_batch:", PER_DEVICE_BATCH, "dtype:", DTYPE)

# ------------------
# Tokenizer loader
# ------------------
sp = spm.SentencePieceProcessor()
sp.load(VOCAB_MODEL)

pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
start_id = sp.piece_to_id("<start>")
end_id = sp.piece_to_id("<end>")
vocab_size = sp.get_piece_size()
print("vocab_size:", vocab_size, "pad_id:", pad_id, "start_id:", start_id, "end_id:", end_id)

# ------------------
# Data pipeline (simple, numpy-based)
# - Reads corpus line-by-line, tokenizes, pads/truncates to SEQ_LEN.
# - Builds a numpy array (N, SEQ_LEN) for inputs and targets (shifted by 1).
# - Shards batches across devices for pmap.
# ------------------
def line_to_ids(line: str, max_len: int = SEQ_LEN):
    ids = sp.encode(line.strip(), out_type=int)
    if len(ids) > max_len - 1:
        ids = ids[: max_len - 1]
    ids = ids + [end_id]
    pad_len = max_len - len(ids)
    ids = ids + [pad_id] * pad_len
    return np.array(ids, dtype=np.int32)

def build_dataset(corpus_path: str, limit: int = LIMIT):
    arr = []
    with open(corpus_path, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i >= limit:
                break
            line = line.strip()
            if not line:
                continue
            arr.append(line_to_ids(line))
    data = np.stack(arr, axis=0)  # (N, SEQ_LEN)
    print("Loaded dataset shape:", data.shape)
    return data

# create inputs and targets
data_np = build_dataset(CORPUS_PATH, LIMIT)
inputs = data_np
targets = np.concatenate([data_np[:,1:], np.full((data_np.shape[0],1), pad_id, dtype=np.int32)], axis=1)

# shuffle and create batches
def create_batch_iter(inputs: np.ndarray, targets: np.ndarray, batch_size: int, rng: np.random.Generator):
    idx = np.arange(inputs.shape[0])
    rng.shuffle(idx)
    for i in range(0, len(idx) - batch_size + 1, batch_size):
        batch_idx = idx[i:i+batch_size]
        x = inputs[batch_idx]
        y = targets[batch_idx]
        yield x, y

# helper to shard numpy batch for pmap: shape (num_devices, per_device, ...)
def shard(xs: np.ndarray):
    return xs.reshape((NUM_DEVICES, -1) + xs.shape[1:])

# ------------------
# Flax model implementation
# ------------------
class SwiGLU(nn.Module):
    d_model: int

    @nn.compact
    def __call__(self, x):
        # project to 2*intermediate, then split
        proj = nn.Dense(self.d_model * 2, dtype=jnp.float32)(x)  # keep proj in float32
        x_val, x_gate = jnp.split(proj, 2, axis=-1)
        out = x_val * nn.silu(x_gate)
        out = nn.Dense(self.d_model, dtype=jnp.float32)(out)
        return out.astype(x.dtype)

class LoU(nn.Module):
    d_model: int
    clip_value: float = 5.0
    eps: float = 1e-6

    @nn.compact
    def __call__(self, x):
        # x: (batch, seq, d)
        x_f32 = x.astype(jnp.float32)
        residual = x_f32

        norm1 = nn.LayerNorm(epsilon=1e-5, dtype=jnp.float32)
        x_norm = norm1(x_f32)

        Q = nn.Dense(self.d_model, dtype=jnp.float32)
        K = nn.Dense(self.d_model, dtype=jnp.float32)
        V = nn.Dense(self.d_model, dtype=jnp.float32)

        q = Q(x_norm)
        k = K(x_norm)
        v = V(x_norm)

        g_q = (jnp.tanh(q) + 1.0) / 2.0
        g_k = (jnp.tanh(k) + 1.0) / 2.0
        score = g_q * g_k  # (b, seq, d)

        alpha_linear = nn.Dense(1, dtype=jnp.float32)
        alpha_dynamic = alpha_linear(x_norm)  # (b, seq, 1)

        # EMA over time: use scan across sequence axis
        # transpose to (seq, batch, d) to scan over time
        score_t = jnp.transpose(score, (1,0,2))
        alpha_t = jnp.transpose(alpha_dynamic, (1,0,2))

        def step(carry, inputs):
            prev_ema = carry
            x_t, a_t = inputs
            new = a_t * x_t + (1.0 - a_t) * prev_ema
            return new, new

        init = score_t[0]
        _, ema_seq = jax.lax.scan(step, init, (score_t[1:], alpha_t[1:]))
        ema_full = jnp.concatenate([init[None, ...], ema_seq], axis=0)  # (seq, batch, d)
        ema = jnp.transpose(ema_full, (1,0,2))  # (batch, seq, d)

        mean_last = jnp.mean(ema, axis=-1, keepdims=True)
        denom = jnp.maximum(mean_last, self.eps)
        score_norm = ema / denom
        score_clipped = jnp.clip(score_norm, -self.clip_value, self.clip_value)

        x_comb = score_clipped * v
        out = x_comb + residual
        out = nn.LayerNorm(epsilon=1e-5, dtype=jnp.float32)(out)
        out = SwiGLU(self.d_model)(out.astype(x.dtype))
        return out.astype(x.dtype)

class Lo(nn.Module):
    d_model: int

    @nn.compact
    def __call__(self, x):
        h = nn.Dense(64, dtype=jnp.float32)(x)
        h = nn.silu(h)
        h = nn.Dense(self.d_model, dtype=jnp.float32)(h)
        out = nn.LayerNorm(epsilon=1e-5, dtype=jnp.float32)(h) + x
        return out.astype(x.dtype)

class Block(nn.Module):
    d_model: int

    @nn.compact
    def __call__(self, x):
        x = LoU(self.d_model)(x)
        x = Lo(self.d_model)(x)
        return x

class ReLM(nn.Module):
    vocab_size: int
    max_seq_len: int
    d_model: int
    n_layers: int
    dtype: Any = jnp.float32

    def setup(self):
        self.token_embed = nn.Embed(self.vocab_size, self.d_model, dtype=self.dtype)
        self.pos_embed = nn.Embed(self.max_seq_len, self.d_model, dtype=self.dtype)
        self.blocks = [Block(self.d_model) for _ in range(self.n_layers)]
        self.ln_f = nn.LayerNorm(epsilon=1e-5, dtype=jnp.float32)

    def __call__(self, x, deterministic=True):
        # x: (batch, seq)
        b, seq = x.shape
        positions = jnp.arange(seq)[None, :]
        x = self.token_embed(x) + self.pos_embed(positions)
        for blk in self.blocks:
            x = blk(x)
        x = self.ln_f(x)
        # tie weights: token embedding matrix
        embedding_matrix = self.token_embed.embedding  # (vocab, d)
        logits = jnp.einsum("bld,vd->blv", x, embedding_matrix)
        return logits.astype(jnp.float32)

# ------------------
# Loss & metrics
# ------------------
def smoothed_cross_entropy(logits, targets, pad_id, eps=0.1):
    # logits: (b, seq, v)
    # targets: (b, seq) int32
    vocab = logits.shape[-1]
    logits = logits.reshape(-1, vocab)
    targets = targets.reshape(-1)
    mask = (targets != pad_id).astype(jnp.float32)
    # one-hot smoothed
    one_hot = jax.nn.one_hot(targets, vocab)
    smooth = (1.0 - eps) * one_hot + eps / float(vocab)
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    loss_per_token = -jnp.sum(smooth * log_probs, axis=-1)
    loss_per_token = loss_per_token * mask
    denom = jnp.sum(mask) + 1e-8
    loss = jnp.sum(loss_per_token) / denom
    return loss

def masked_perplexity_from_logits(logits, targets, pad_id, eps=0.1):
    vocab = logits.shape[-1]
    logits = logits.reshape(-1, vocab)
    targets = targets.reshape(-1)
    mask = (targets != pad_id).astype(jnp.float32)
    one_hot = jax.nn.one_hot(targets, vocab)
    smooth = (1.0 - eps) * one_hot + eps / float(vocab)
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    loss_per_token = -jnp.sum(smooth * log_probs, axis=-1) * mask
    mean_loss = jnp.sum(loss_per_token) / (jnp.sum(mask) + 1e-8)
    return jnp.exp(mean_loss)

# ------------------
# Training state
# ------------------
class TrainState(train_state.TrainState):
    pass

def create_train_state(rng, model, learning_rate):
    params = model.init(rng, jnp.zeros((1, SEQ_LEN), dtype=jnp.int32))["params"]
    tx = optax.chain(
        optax.clip_by_global_norm(1.0),
        optax.adamw(learning_rate=learning_rate, b1=0.9, b2=0.95, eps=1e-8)
    )
    return TrainState.create(apply_fn=model.apply, params=params, tx=tx)

# ------------------
# pmap'd step functions
# ------------------
@partial(jax.pmap, axis_name="batch")
def train_step(state, batch_x, batch_y, rng):
    def loss_fn(params):
        logits = state.apply_fn({"params": params}, batch_x, deterministic=False)
        loss = smoothed_cross_entropy(logits, batch_y, pad_id)
        return loss, logits

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, logits), grads = grad_fn(state.params)
    grads = jax.lax.pmean(grads, axis_name="batch")
    new_state = state.apply_gradients(grads=grads)
    # metrics
    ppl = masked_perplexity_from_logits(logits, batch_y, pad_id)
    metrics = {"loss": loss, "ppl": ppl}
    metrics = jax.lax.pmean(metrics, axis_name="batch")
    return new_state, metrics

@partial(jax.pmap, axis_name="batch")
def eval_step(state, batch_x, batch_y):
    logits = state.apply_fn({"params": state.params}, batch_x, deterministic=True)
    loss = smoothed_cross_entropy(logits, batch_y, pad_id)
    ppl = masked_perplexity_from_logits(logits, batch_y, pad_id)
    metrics = {"loss": loss, "ppl": ppl}
    metrics = jax.lax.pmean(metrics, axis_name="batch")
    return metrics

# ------------------
# Training loop
# ------------------
rng = random.PRNGKey(SEED)
rng, init_rng = random.split(rng)
model = ReLM(vocab_size=vocab_size, max_seq_len=SEQ_LEN, d_model=512, n_layers=9, dtype=DTYPE)
state = create_train_state(init_rng, model, LEARNING_RATE)

# replicate to devices
state = jax.device_put_replicated(state, jax.local_devices())

print("Starting training...")

global_step = 0
for epoch in range(EPOCHS):
    print(f"Epoch {epoch+1}/{EPOCHS}")
    np_rng = np.random.default_rng(SEED + epoch)
    batch_iter = create_batch_iter(inputs, targets, GLOBAL_BATCH, np_rng)
    pbar = tqdm.tqdm(batch_iter, total= max(1, inputs.shape[0] // GLOBAL_BATCH))

    for batch_x, batch_y in pbar:
        # shard
        batch_x = shard(batch_x)
        batch_y = shard(batch_y)
        rng, step_rng = random.split(rng)
        # make per-device rngs
        step_rngs = random.split(step_rng, NUM_DEVICES)
        state, metrics = train_step(state, batch_x, batch_y, step_rngs)
        # metrics are per-device; take first replica
        m = jax.tree_util.tree_map(lambda x: x[0], metrics)
        pbar.set_postfix(loss=float(m["loss"]), ppl=float(m["ppl"]))
        global_step += 1

# ------------------
# Save params
# ------------------
save_dir = "./checkpoints"
os.makedirs(save_dir, exist_ok=True)
# save using flax.serialization via checkpoints
checkpoints.save_checkpoint(save_dir, jax.tree_map(lambda x: np.array(x), state), step=global_step, keep=3)
print("Saved checkpoint to", save_dir)

# ------------------
# Sampling (top-p) - single-device (CPU) sampling for simplicity
# ------------------
import math

def top_p_sample_logits(rng, logits, p=0.9, temperature=1.0):
    # logits: (vocab,)
    probs = jax.nn.softmax(logits / temperature)
    # convert to numpy for sorting (ok for single token)
    probs_np = np.array(probs)
    sorted_idx = np.argsort(probs_np)[::-1]
    sorted_probs = probs_np[sorted_idx]
    cum = np.cumsum(sorted_probs)
    cutoff = np.searchsorted(cum, p)
    top_idx = sorted_idx[: cutoff + 1]
    top_probs = sorted_probs[: cutoff + 1]
    top_probs = top_probs / top_probs.sum()
    # sample
    next_token = np.random.choice(top_idx, p=top_probs)
    return int(next_token)

def generate_text(state, prompt: str, max_gen=256, p=0.9, temperature=0.8, min_len=20):
    # load params from replicated state (take first replica)
    params = jax.tree_map(lambda x: np.array(x[0]), state.params)
    tokens = sp.encode("<start> " + prompt, out_type=int)
    generated = tokens.copy()
    for step in range(max_gen):
        cur = generated[-SEQ_LEN:]
        if len(cur) < SEQ_LEN:
            cur = cur + [pad_id] * (SEQ_LEN - len(cur))
        x = np.array([cur], dtype=np.int32)
        logits = model.apply({"params": params}, x, deterministic=True)  # (1, seq, vocab)
        logits = np.array(logits[0, len(generated)-1 if len(generated)-1 < SEQ_LEN else SEQ_LEN-1])
        # penalize end/pad a bit
        logits[end_id] -= 5.0
        logits[pad_id] -= 10.0
        next_id = top_p_sample_logits(None, logits, p=p, temperature=temperature)
        generated.append(next_id)
        if next_id == end_id and len(generated) >= min_len:
            break
    return sp.decode(generated)

# quick generate
print("\n\n===== 생성 결과 =====")
print(generate_text(state, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))