"""Train the six-score head on ELLIPSE, on one Modal A10G.

    MODAL_PROFILE=labs3 modal run scripts/train_scores.py --tag scores-v1
    MODAL_PROFILE=labs3 modal run scripts/train_scores.py --tag smoke --limit 64 --epochs 1

Metric: MCRMSE over the six traits on the ELLIPSE test split (the Kaggle winners' single models: ~0.45 large, ~0.46-0.47
base). Weights and the log go to /data/runs/<tag> and HF nonsodev/essay-coach runs/<tag>.
"""
import os

import modal

app = modal.App("essay-coach-train")
HF_SECRET = True  # flip to True once `modal secret create huggingface HF_TOKEN=...` exists on this workspace
SECRETS = [modal.Secret.from_name("huggingface")] if HF_SECRET else []
vol = modal.Volume.from_name("essay-data")
cache = modal.Volume.from_name("essay-hf-cache", create_if_missing=True)
image = (modal.Image.debian_slim(python_version="3.11")
         .pip_install("torch==2.4.1", index_url="https://download.pytorch.org/whl/cu124")
         .pip_install("transformers==4.46.3", "sentencepiece", "protobuf", "tiktoken", "safetensors", "pandas", "numpy<2.3",
                      "scikit-learn", "huggingface_hub")
         .add_local_python_source("coach"))
HF_REPO = "nonsodev/essay-coach"


@app.function(image=image, gpu="A10G", volumes={"/data": vol, "/root/.cache/huggingface": cache}, timeout=4 * 3600,
              secrets=SECRETS, cpu=4, memory=16384)
def train(tag: str, epochs: int, bs: int, accum: int, lr: float, max_len: int, limit: int, seed: int):
    import json
    import os
    import time

    import numpy as np
    import torch
    from huggingface_hub import HfApi
    from safetensors.torch import save_file

    from coach.data import TRAITS, ellipse, with_breaks
    from coach.model import ScoreModel, param_groups, tokenizer

    torch.manual_seed(seed)
    np.random.seed(seed)
    run = f"/data/runs/{tag}"
    os.makedirs(run, exist_ok=True)
    tr, te = ellipse()
    if limit:
        tr, te = tr.sample(limit, random_state=seed), te.sample(limit, random_state=seed)
    print(f"train {len(tr)} test {len(te)}", flush=True)
    tok = tokenizer()

    def encode(df):
        enc = tok([with_breaks(t) for t in df.text], truncation=True, max_length=max_len)
        y = df[TRAITS].values.astype(np.float32)
        return [(enc["input_ids"][i], y[i]) for i in range(len(df))]

    def batches(items, size, shuffle):
        idx = np.random.permutation(len(items)) if shuffle else np.arange(len(items))
        for s in range(0, len(idx), size):
            chunk = [items[i] for i in idx[s:s + size]]
            L = max(len(ids) for ids, _ in chunk)
            ids = torch.full((len(chunk), L), tok.pad_token_id)
            att = torch.zeros((len(chunk), L), dtype=torch.long)
            for j, (seq, _) in enumerate(chunk):
                ids[j, :len(seq)] = torch.tensor(seq)
                att[j, :len(seq)] = 1
            yield ids.cuda(), att.cuda(), torch.tensor(np.stack([y for _, y in chunk])).cuda()

    train_items, test_items = encode(tr), encode(te)
    net = ScoreModel(n_tokens=len(tok)).cuda()
    opt = torch.optim.AdamW(param_groups(net, lr, head_lr=1e-3))
    steps = epochs * ((len(train_items) + bs - 1) // bs) // accum
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=[g["lr"] for g in opt.param_groups], total_steps=max(steps, 1),
                                                pct_start=0.1, anneal_strategy="linear")

    @torch.no_grad()
    def evaluate():
        net.eval()
        preds, ys = [], []
        for ids, att, y in batches(test_items, 16, False):
            with torch.autocast("cuda", torch.bfloat16):
                preds.append(net(ids, att).float().cpu())
            ys.append(y.cpu())
        net.train()
        p, y = torch.cat(preds).numpy(), torch.cat(ys).numpy()
        rmse = np.sqrt(((p - y) ** 2).mean(0))
        return float(rmse.mean()), {t: round(float(r), 4) for t, r in zip(TRAITS, rmse)}, p

    hf = bool(os.environ.get("HF_TOKEN"))
    if hf:
        api = HfApi()
        api.create_repo(HF_REPO, exist_ok=True)
    log, best, t0, step = [], 9.0, time.time(), 0
    net.train()
    for ep in range(epochs):
        tot, n = 0.0, 0
        for i, (ids, att, y) in enumerate(batches(train_items, bs, True)):
            with torch.autocast("cuda", torch.bfloat16):
                loss = torch.nn.functional.mse_loss(net(ids, att).float(), y) / accum
            loss.backward()
            tot, n = tot + loss.item() * accum, n + 1
            if (i + 1) % accum == 0:
                torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0)
                opt.step()
                opt.zero_grad(set_to_none=True)
                if step < steps - 1:
                    sched.step()
                step += 1
            if n % 100 == 0:
                print(f"epoch {ep} batch {n} loss {tot/n:.4f} {time.time()-t0:.0f}s", flush=True)
        mcrmse, per, p = evaluate()
        row = {"epoch": ep, "train_mse": round(tot / n, 4), "mcrmse": round(mcrmse, 4), **per, "t": round(time.time() - t0)}
        log.append(row)
        print(row, flush=True)
        if mcrmse < best:
            best = mcrmse
            save_file({k: v.detach().cpu().contiguous() for k, v in net.state_dict().items()}, f"{run}/model.safetensors")
            te.assign(**{f"pred_{t}": p[:, j] for j, t in enumerate(TRAITS)}).to_csv(f"{run}/test_preds.csv", index=False)
        with open(f"{run}/log.jsonl", "w") as f:
            f.write("\n".join(json.dumps(r) for r in log) + "\n")
        vol.commit()
    tok.save_pretrained(f"{run}/tokenizer")
    json.dump({"backbone": net.backbone.config._name_or_path, "traits": TRAITS, "max_len": max_len, "best_mcrmse": best},
              open(f"{run}/config.json", "w"))
    vol.commit()
    if hf:
        api.upload_folder(folder_path=run, repo_id=HF_REPO, path_in_repo=f"runs/{tag}", commit_message=f"{tag} mcrmse {best:.4f}")
    return {"tag": tag, "best_mcrmse": round(best, 4), "minutes": round((time.time() - t0) / 60, 1)}


@app.local_entrypoint()
def main(tag: str = "scores-v1", epochs: int = 4, bs: int = 8, accum: int = 1, lr: float = 2e-5, max_len: int = 1024,
         limit: int = 0, seed: int = 0):
    print(train.remote(tag, epochs, bs, accum, lr, max_len, limit, seed))