"""PC-SHO-DLM: LOCAL PC Training on A100 — Single mode, fast."""
import json, os, sys, threading, time
import gradio as gr
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater, count_parameters

class CharDS(Dataset):
    def __init__(self, text, seq_len):
        self.seq_len = seq_len
        self.data = torch.tensor([min(b+1,256) for b in text.encode("utf-8")], dtype=torch.long)
        self.n = max(1, (len(self.data)-seq_len)//seq_len)
    def __len__(self): return self.n
    def __getitem__(self, i):
        return {"input_ids": self.data[i*self.seq_len:(i+1)*self.seq_len]}

LOG = []
STATUS = "Idle"

def train():
    global STATUS, LOG
    LOG = ["Starting LOCAL PC training on A100..."]
    device = "cuda" if torch.cuda.is_available() else "cpu"
    LOG.append(f"Device: {device}")
    if device == "cuda":
        LOG.append(f"GPU: {torch.cuda.get_device_name()}")
        LOG.append(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")

    config = PCSHOConfig(
        vocab_size=257, max_seq_len=512, d_model=512, n_heads=8,
        n_layers=12, d_ff=2048, n_diffusion_steps=128, n_settling_steps=6,
        mask_token_id=0, dropout=0.1, feedback_rank=128,
    )
    LOG.append(f"Model: {count_parameters(PCSHODLM(config)):,} params")

    STATUS = "Loading data..."
    from datasets import load_dataset
    ds = load_dataset("wikitext","wikitext-103-raw-v1",split="train")
    dv = load_dataset("wikitext","wikitext-103-raw-v1",split="validation")
    tt = "\n".join([r["text"] for r in ds if r["text"].strip()])[:100_000_000]
    vt = "\n".join([r["text"] for r in dv if r["text"].strip()])
    train_ds, val_ds = CharDS(tt,config.max_seq_len), CharDS(vt,config.max_seq_len)
    LOG.append(f"Data: {len(tt):,} train, {len(vt):,} val chars")

    model = PCSHODLM(config).to(device)
    model.train()
    updater = LocalParameterUpdater(model, lr_forward=3e-4, lr_feedback=3e-4, lr_readout=3e-4, lr_precision=3e-5)
    tl = DataLoader(train_ds, batch_size=64, shuffle=True, drop_last=True, num_workers=4, pin_memory=True)
    vl = DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=2, pin_memory=True)

    STATUS = "Training LOCAL PC..."
    LOG.append(f"\n{'='*60}\nMode: LOCAL PC | Params: {count_parameters(model):,}\n{'='*60}")
    step, start, max_steps = 0, time.time(), 10000
    log_data = {"steps":[],"losses":[],"energies":[]}

    while step < max_steps:
        for batch in tl:
            if step >= max_steps: break
            x_0 = batch["input_ids"].to(device)
            result = updater.step({"input_ids": x_0})
            step += 1
            loss = result.get("loss",0); energies = result.get("energies",[])
            if step % 100 == 0:
                elapsed = time.time()-start
                tps = step*64*config.max_seq_len/elapsed
                e = f"{energies[-1]:.0f}" if energies else "N/A"
                msg = f"[local] Step {step:6d} | Loss: {loss:.4f} | Energy: {e} | Tok/s: {tps:.0f} | {elapsed:.0f}s"
                LOG.append(msg); log_data["steps"].append(step); log_data["losses"].append(loss); log_data["energies"].append(energies[-1] if energies else 0)
            if step % 2000 == 0:
                model.eval()
                tl2,tt2=0.0,0
                with torch.no_grad():
                    for vb in vl:
                        vx=vb["input_ids"].to(device); vo=model(vx); nm=vo["mask"].sum().item()
                        if nm>0: tl2+=vo["loss"].item()*nm; tt2+=nm
                        if tt2>100000: break
                vl2=tl2/max(1,tt2); LOG.append(f"  --> [local] Val loss: {vl2:.4f}"); log_data.setdefault("val",[]).append((step,vl2))
                model.train()

    model.eval()
    tl2,tt2=0.0,0
    with torch.no_grad():
        for vb in vl:
            vx=vb["input_ids"].to(device); vo=model(vx); nm=vo["mask"].sum().item()
            if nm>0: tl2+=vo["loss"].item()*nm; tt2+=nm
            if tt2>200000: break
    fv=tl2/max(1,tt2); log_data["final_val"]=fv
    LOG.append(f"DONE: {step} steps in {time.time()-start:.0f}s | Final val: {fv:.4f}")
    STATUS = "Complete!"

    os.makedirs("results",exist_ok=True)
    torch.save({"model":model.state_dict(),"config":config,"log":log_data},"results/local_pc_10k.pt")
    with open("results/local_pc_log.json","w") as f: json.dump(log_data,f)
    try:
        from huggingface_hub import HfApi
        HfApi().upload_folder(folder_path="results",repo_id="zotowata/pc-sho-dlm-local",repo_type="space",path_in_repo="results")
        LOG.append("Results uploaded!")
    except: pass

_thread = None
def start():
    global _thread
    if _thread and _thread.is_alive(): return "Already running!"
    _thread = threading.Thread(target=train, daemon=True); _thread.start()
    return "LOCAL PC training started on A100!"

with gr.Blocks(title="PC-SHO-DLM Local PC") as demo:
    gr.Markdown("# PC-SHO-DLM: LOCAL PC Training (A100)")
    with gr.Row():
        btn = gr.Button("Start Training", variant="primary")
        st = gr.Textbox(label="Status", value="Idle")
    log = gr.Textbox(label="Log", lines=25, max_lines=50)
    btn.click(start, outputs=st)
    refresh = gr.Button("Refresh"); refresh.click(lambda: "\n".join(LOG[-50:]), outputs=log); refresh.click(lambda: STATUS, outputs=st)
    timer = gr.Timer(5); timer.tick(lambda: "\n".join(LOG[-50:]), outputs=log); timer.tick(lambda: STATUS, outputs=st)

demo.launch()