"""TensorPEFT live demo: tensor gradient-pursuit adapters vs. matched-budget PEFT baselines.

GPT-2 (124M, English) is fine-tuned on Bengali news with the same number of trainable
parameters for every method; validation cross-entropy (nats/token) decides.
"""
import copy
import glob
import io
import json
import math
import os
import random
import re
import time
import unicodedata
import urllib.request
import zipfile

import gradio as gr
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402

from baselines import IA3, LNTune, LoRA, Prompt, _Adapter  # noqa: E402
from tensorpeft import TensorPEFTConfig, get_tensorpeft_model  # noqa: E402

try:
    import spaces
    GPU = spaces.GPU
except Exception:  # running outside a ZeroGPU Space
    def GPU(fn=None, **kw):
        return fn if fn is not None else (lambda f: f)

MODEL_ID = os.environ.get("TP_MODEL", "openai-community/gpt2")
DATA_URL = os.environ.get(
    "TP_DATA_URL",
    "https://raw.githubusercontent.com/soham96/Bengali_news_classifier/"
    "4c2a51b51fc56208426ef19368bfcb08acc9d0c4/data/data.zip")
CACHE = os.environ.get("TP_CACHE", "data_cache")
SEQ, BATCH, VAL_SEQS = 128, 8, 64

# ------------------------------------------------------------------ data (built once at start-up)
tok = AutoTokenizer.from_pretrained(MODEL_ID)


def build_data():
    os.makedirs(CACHE, exist_ok=True)
    tr, va = os.path.join(CACHE, "train.bin"), os.path.join(CACHE, "val.bin")
    if os.path.exists(tr) and os.path.exists(va):
        return
    src = DATA_URL
    raw = open(src, "rb").read() if os.path.exists(src) else urllib.request.urlopen(src).read()
    docs = []
    with zipfile.ZipFile(io.BytesIO(raw)) as z:
        for name in sorted(n for n in z.namelist() if n.endswith("_articles.txt") and "__MACOSX" not in n):
            for a in json.loads(z.read(name).decode("utf-8"))["articles"]:
                t = unicodedata.normalize("NFC", (a.get("title", "") + "। " + a.get("body", "")).strip())
                t = re.sub(r"\s+", " ", t)
                if len(t) > 200:
                    docs.append(t)
    random.Random(0).shuffle(docs)

    def enc(ds, cap):
        ids = []
        for d in ds:
            ids += tok(d).input_ids + [tok.eos_token_id]
            if len(ids) >= cap:
                break
        return np.array(ids[:cap], dtype=np.uint16)

    enc(docs[:400], 200_000).tofile(va)
    enc(docs[400:], 3_000_000).tofile(tr)


build_data()
TRAIN = np.fromfile(os.path.join(CACHE, "train.bin"), dtype=np.uint16)
VAL = torch.from_numpy(np.fromfile(os.path.join(CACHE, "val.bin"), dtype=np.uint16)[: VAL_SEQS * (SEQ + 1)]
                       .astype(np.int64).reshape(VAL_SEQS, SEQ + 1))
BASE = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32).eval()
for _p in BASE.parameters():
    _p.requires_grad_(False)


def sampler(seed):
    rng = np.random.default_rng(seed)
    while True:
        ix = rng.integers(0, len(TRAIN) - SEQ - 1, BATCH)
        yield torch.from_numpy(np.stack([TRAIN[i:i + SEQ + 1] for i in ix]).astype(np.int64))


# ------------------------------------------------------------------ methods per budget
class TP(_Adapter):
    lr, warmup = 3e-3, 10

    def __init__(self, model, **cfg):
        super().__init__()
        self.model = [model]
        self.tp = get_tensorpeft_model(model, TensorPEFTConfig(**cfg))

    def setup(self, probe_batches):
        self.tp.init_atoms(probe_batches, verbose=False)

    def trainable(self):
        return [p for p in self.tp.parameters() if p.requires_grad]

    def mark_ready(self):
        for g in self.tp.groups:
            g.initialized.fill_(True)


ATTN = ["c_attn.q", "c_attn.k", "c_attn.v", "attn.c_proj"]
BUDGETS = {
    "≈1.5k parameters": [
        ("TensorPEFT · 1 attention atom", lambda m: TP(m, target_modules=ATTN, atoms_per_group=1)),
        ("Prompt tuning · 1 token", lambda m: Prompt(m, 1)),
        ("LayerNorm tuning · final LN", lambda m: LNTune(m, layers=[], final=True)),
    ],
    "≈9k parameters": [
        ("TensorPEFT · 3 atoms", lambda m: TP(m, atoms_per_group=1)),
        ("Prompt tuning · 12 tokens", lambda m: Prompt(m, 12)),
        ("LoRA r=1 · q,v of last 3 layers", lambda m: LoRA(m, 1, layers=[9, 10, 11])),
        ("(IA)³ · values", lambda m: IA3(m, kv=("c_attn.v",), ff=False)),
        ("LayerNorm tuning · last 3 layers", lambda m: LNTune(m, layers=[9, 10, 11], final=False)),
    ],
    "≈55k parameters": [
        ("TensorPEFT · 18 atoms", lambda m: TP(m, atoms_per_group=6)),
        ("(IA)³ · keys, values, MLP", lambda m: IA3(m)),
    ],
}
COLORS = ["#7c3aed", "#0ea5e9", "#f59e0b", "#10b981", "#ef4444"]


def ce(adapter, x):
    return F.cross_entropy(adapter.logits(x[:, :-1]).flatten(0, 1), x[:, 1:].flatten())


@torch.no_grad()
def evaluate(adapter, dev):
    V = VAL.to(dev)
    return sum(float(ce(adapter, V[i:i + 16])) * len(V[i:i + 16]) for i in range(0, len(V), 16)) / len(V)


def _duration(budget, steps, seed):
    return int(min(280, 30 + len(BUDGETS[budget]) * steps * 0.06))


@GPU(duration=_duration)
def race_gpu(budget, steps, seed):
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    if dev == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
    out = []
    plain = _Adapter()
    plain.model = [copy.deepcopy(BASE).to(dev)]
    base_val = evaluate(plain, dev)
    del plain
    for mi, (name, build) in enumerate(BUDGETS[budget]):
        torch.manual_seed(seed)
        model = copy.deepcopy(BASE).to(dev).eval()
        ad = build(model)
        t0 = time.time()
        if isinstance(ad, TP):
            ad.setup([b.to(dev) for b, _ in zip(sampler(seed + 12345), range(4))])
        params = ad.trainable()
        opt = torch.optim.AdamW(params, lr=ad.lr, weight_decay=0.0)
        warm = getattr(ad, "warmup", 0)
        batches = sampler(seed)
        curve = [(0, evaluate(ad, dev))]
        for step in range(1, steps + 1):
            if warm:
                for g in opt.param_groups:
                    g["lr"] = ad.lr * min(1.0, step / warm)
            loss = ce(ad, next(batches).to(dev))
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            if step % 25 == 0 or step == steps:
                curve.append((step, evaluate(ad, dev)))
        out.append(dict(name=name, n_params=ad.n_params(), curve=curve, final=curve[-1][1],
                        seconds=time.time() - t0,
                        state=[p.detach().float().cpu().clone() for p in params]))
        del model, ad, opt
    return dict(budget=budget, steps=steps, seed=seed, base=base_val, runs=out)


def run_race(budget, steps, seed):
    res = race_gpu(budget, int(steps), int(seed))
    fig, ax = plt.subplots(figsize=(8, 4.6))
    for i, r in enumerate(res["runs"]):
        xs, ys = zip(*r["curve"])
        ax.plot(xs, ys, color=COLORS[i % len(COLORS)], lw=2.6 if r["name"].startswith("TensorPEFT") else 1.6,
                label=f"{r['name']} ({r['n_params']:,} params)")
    ax.axhline(res["base"], color="#9ca3af", ls="--", lw=1, label=f"GPT-2 before fine-tuning ({res['base']:.3f})")
    ax.set_xlabel("training step")
    ax.set_ylabel("validation loss (nats/token)  ↓ better")
    ax.set_title(f"Bengali fine-tuning of GPT-2 at {res['budget']}")
    ax.grid(alpha=0.25)
    ax.legend(fontsize=8, loc="upper right")
    fig.tight_layout()
    ranked = sorted(res["runs"], key=lambda r: r["final"])
    best = ranked[0]
    rows = [[i + 1, r["name"], f"{r['n_params']:,}", f"{r['final']:.4f}", f"{r['final'] - best['final']:+.4f}",
             f"{r['seconds']:.0f}s"] for i, r in enumerate(ranked)]
    tp = next(r for r in res["runs"] if r["name"].startswith("TensorPEFT"))
    others = [r for r in res["runs"] if r is not tp]
    nxt = min(others, key=lambda r: r["final"])
    verdict = (f"### {'🏆 TensorPEFT wins' if tp is best else 'TensorPEFT does not win this run'}\n"
               f"TensorPEFT **{tp['final']:.4f}** vs best baseline ({nxt['name']}) **{nxt['final']:.4f}** "
               f"→ difference **{tp['final'] - nxt['final']:+.4f}** nats/token after {res['steps']} steps (seed {res['seed']}).")
    store = {"budget": res["budget"], "runs": [{"name": r["name"], "state": r["state"]} for r in res["runs"]]}
    return fig, rows, verdict, store


@GPU(duration=90)
def generate_gpu(store, prompt, n_new, temperature, seed):
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    ids0 = tok(prompt, return_tensors="pt").input_ids.to(dev)
    builders = dict(BUDGETS[store["budget"]]) if store else {}
    entries = [("GPT-2 before fine-tuning", None, None)] + \
              [(r["name"], builders[r["name"]], r["state"]) for r in (store["runs"] if store else [])]
    outs = []
    for name, build, state in entries:
        model = copy.deepcopy(BASE).to(dev).eval()
        if build is None:
            ad = _Adapter()
            ad.model = [model]
        else:
            ad = build(model)
            if isinstance(ad, TP):
                ad.mark_ready()
            with torch.no_grad():
                for p, s in zip(ad.trainable(), state):
                    p.copy_(s.to(p.device, p.dtype))
        g = torch.Generator(device=dev).manual_seed(int(seed))
        ids = ids0.clone()
        with torch.no_grad():
            for _ in range(int(n_new)):
                logits = ad.logits(ids[:, -512:])[:, -1].float() / max(float(temperature), 1e-3)
                top = torch.topk(logits, 50)
                probs = torch.softmax(top.values, -1)
                nxt = top.indices.gather(-1, torch.multinomial(probs, 1, generator=g))
                ids = torch.cat([ids, nxt], dim=1)
        outs.append((name, tok.decode(ids[0, ids0.shape[1]:], skip_special_tokens=True)))
        del model, ad
    return outs


def run_generate(store, prompt, n_new, temperature, seed):
    if not prompt.strip():
        return "Enter a Bengali prompt."
    outs = generate_gpu(store, prompt, n_new, temperature, seed)
    note = "" if store else "*Run a race first to compare the fine-tuned models; showing the base model only.*\n\n"
    return note + "\n\n".join(f"**{n}**\n\n> {prompt}{t}" for n, t in outs)


# ------------------------------------------------------------------ benchmark results (CPU, 300 steps)
RESULTS_MD = """
### Three seeds (mean ± std of validation loss; lower is better)

| budget | TensorPEFT | best baseline | difference | seeds won | p (paired t) |
|---|---:|---|---:|:---:|---:|
| ≈9k | **1.3245 ± 0.004** | prompt tuning 1.3684 ± 0.010 | **−0.044** | 3/3 | 0.008 |
| ≈37k | 1.2558 ± 0.010 | LayerNorm tuning (38.4k) 1.2504 ± 0.003 | +0.005 | 1/3 | 0.37 (tie) |
| ≈55k | **1.2368 ± 0.017** | (IA)³ 1.2880 ± 0.003 | **−0.051** | 3/3 | 0.026 |

### Every method at every budget (seed 0)

| budget | TensorPEFT | LoRA | QLoRA | VeRA | (IA)³ | prompt tuning | LayerNorm tuning | BitFit |
|---|---:|---:|---:|---:|---:|---:|---:|---:|
| ≈1.5k | **1.461** | – | – | – | – | 1.493 (768) | 1.600 | – |
| ≈9k | **1.327** | 1.430 | – | – | 1.432 | 1.370 | 1.407 | – |
| ≈37k | 1.255 | 1.318 | 1.335 | 1.393 | – | 1.327 | **1.253** | – |
| ≈55k | **1.231** | – | – | – | 1.289 | – | – | – |
| ≈83k | – | – | – | – | – | – | – | 1.326 |

GPT-2 124M before fine-tuning: 1.708. 300 AdamW steps, batch 8 × 128 tokens, identical data for all methods,
baseline learning rates from their papers. “–”: the method cannot be configured at that budget on GPT-2 or was
not run there (LoRA's smallest all-layer setting is 37k). At larger budgets (≥ 115k parameters) LoRA was
better in our tests — TensorPEFT is a method for very small budgets.
"""


def results_figure():
    pts = {
        "TensorPEFT": [(1553, 1.461), (9259, 1.327), (37036, 1.255), (55554, 1.231)],
        "LoRA / QLoRA": [(9216, 1.430), (36864, 1.318), (36864, 1.335)],
        "Prompt tuning": [(768, 1.493), (9216, 1.370), (36864, 1.327)],
        "LayerNorm tuning": [(1536, 1.600), (9216, 1.407), (38400, 1.253)],
        "(IA)³": [(9216, 1.432), (55296, 1.289)],
        "VeRA": [(24576, 1.393)],
        "BitFit": [(82944, 1.326)],
    }
    fig, ax = plt.subplots(figsize=(8, 4.6))
    for i, (k, v) in enumerate(pts.items()):
        xs, ys = zip(*v)
        if k == "TensorPEFT":
            ax.plot(xs, ys, "-o", color=COLORS[0], lw=2.6, ms=7, label=k, zorder=5)
        else:
            ax.scatter(xs, ys, s=45, color=["#0ea5e9", "#f59e0b", "#10b981", "#ef4444", "#64748b", "#a16207"][i - 1],
                       label=k)
    ax.set_xscale("log")
    ax.set_xlabel("trainable parameters (log scale)")
    ax.set_ylabel("validation loss (nats/token)  ↓ better")
    ax.set_title("GPT-2 → Bengali: loss vs. trainable parameters (seed 0)")
    ax.grid(alpha=0.25, which="both")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


# ------------------------------------------------------------------ UI
INTRO = """
# TensorPEFT · tensor gradient-pursuit fine-tuning
LoRA gives every weight matrix its own low-rank update. **TensorPEFT** stacks all equal-shaped matrices of the
network into one tensor (layers × {q,k,v,o} × d_in × d_out) and learns a few **rank-one CP atoms** of it, chosen from
one gradient probe before training. Each atom touches every layer for only L + P + d_in + d_out + 1 parameters.

Here GPT-2 (English) is taught **Bengali** with the *same number of trainable parameters* for every method.
`pip install git+https://huggingface.co/Angshul/TensorPEFT`
"""

with gr.Blocks(title="TensorPEFT") as demo:
    gr.Markdown(INTRO)
    store = gr.State(None)
    with gr.Tab("Live race"):
        with gr.Row():
            budget = gr.Radio(list(BUDGETS), value="≈9k parameters", label="Trainable-parameter budget")
            steps = gr.Number(value=300, visible=False)
            seed = gr.Number(value=0, precision=0, label="Seed")
        gr.Markdown("Every method trains for **300 AdamW steps** (batch 8 × 128 tokens) on identical batches — the setting of the benchmark. "
                    "TensorPEFT starts from zero coefficients, so it trails in the first few dozen steps and overtakes later.")
        go = gr.Button("Run the race", variant="primary")
        verdict = gr.Markdown()
        plot = gr.Plot()
        table = gr.Dataframe(headers=["rank", "method", "params", "val loss", "vs best", "time"], interactive=False)
        go.click(run_race, [budget, steps, seed], [plot, table, verdict, store])
    with gr.Tab("Generate"):
        gr.Markdown("Continuations from the models trained in your last race (top-k 50 sampling).")
        prompt = gr.Textbox(value="কলকাতা শহরে আজ", label="Bengali prompt")
        with gr.Row():
            n_new = gr.Slider(20, 120, value=60, step=10, label="New tokens")
            temp = gr.Slider(0.3, 1.2, value=0.8, step=0.1, label="Temperature")
            gseed = gr.Number(value=0, precision=0, label="Seed")
        gen = gr.Button("Generate", variant="primary")
        gen_out = gr.Markdown()
        gen.click(run_generate, [store, prompt, n_new, temp, gseed], gen_out)
    with gr.Tab("Benchmark results"):
        gr.Plot(value=results_figure())
        gr.Markdown(RESULTS_MD)

if __name__ == "__main__":
    demo.launch()