import os, random, hashlib, json, time, subprocess, pathlib
import tinker, datasets, numpy as np
from tinker import types
from tinker_cookbook import renderers, tokenizer_utils
from tinker_cookbook.supervised.data import conversation_to_datum

MODEL = "Qwen/Qwen3.5-2B"; BS = 32; MAXLEN = 2048; RANK = 32
FORK_URL = open(pathlib.Path(__file__).with_name("fork_url.txt")).read().strip()
VOLUME = "run_739abc9a8b-checkpoints"; ENV = "research-exchange-runs"

def service():
    return tinker.ServiceClient(
        base_url=FORK_URL,
        api_key=open(pathlib.Path(__file__).with_name(".fork_api_key")).read().strip(),
        default_headers_overrides={"Modal-Key": os.environ["RUNS_PROXY_TOKEN_ID"],
                                   "Modal-Secret": os.environ["RUNS_PROXY_TOKEN_SECRET"]})

_tok = _rend = None
def renderer():
    global _tok, _rend
    if _rend is None:
        _tok = tokenizer_utils.get_tokenizer(MODEL)
        name = renderers.get_recommended_renderer_name(MODEL) if hasattr(renderers, "get_recommended_renderer_name") else "qwen3"
        _rend = renderers.get_renderer(name, _tok)
    return _rend

def splits():
    ds = datasets.load_dataset("HuggingFaceH4/no_robots")
    tr = list(ds["train"]); random.Random(1234).shuffle(tr)
    return tr[:500], tr[500:], list(ds["test"])

def train_order(train, seed):
    t = list(train); random.Random(seed).shuffle(t); return t

def mk(r):
    return conversation_to_datum(r["messages"], renderer(), MAXLEN)

def batch(order, step):
    rows = order[step * BS:(step + 1) * BS]
    return [mk(r) for r in rows]

def batch_hash(data):
    h = hashlib.sha256()
    for d in data:
        h.update(np.asarray(d.model_input.to_ints(), dtype=np.int64).tobytes())
        for k in ("target_tokens", "weights"):
            h.update(np.asarray(d.loss_fn_inputs[k].tolist(), dtype=np.float64).tobytes())
    return h.hexdigest()

def opt_kw(opt):
    return {"optimizer": types.DimuonOptimizerConfig()} if opt == "dimuon" else {}

def step_params(opt, lr):
    if opt == "dimuon":
        return types.DimuonParams(learning_rate=lr, beta1=0.9, grad_clip_norm=0.0)
    return tinker.AdamParams(learning_rate=lr, beta1=0.9, beta2=0.95, eps=1e-8)

def token_nll(fb_out, data):
    s = n = 0.0
    for o, d in zip(fb_out.loss_fn_outputs, data):
        lp = np.asarray(o["logprobs"].tolist(), dtype=np.float64)
        w = np.asarray(d.loss_fn_inputs["weights"].tolist(), dtype=np.float64)
        m = w > 0; s += float(-lp[m].sum()); n += float(m.sum())
    return s, n

def eval_nll(tc, data, chunk=32):
    futs = [tc.forward(data[i:i + chunk], loss_fn="cross_entropy") for i in range(0, len(data), chunk)]
    s = n = 0.0; per = []
    for i, f in zip(range(0, len(data), chunk), futs):
        out = f.result()
        for o, d in zip(out.loss_fn_outputs, data[i:i + chunk]):
            lp = np.asarray(o["logprobs"].tolist(), dtype=np.float64)
            w = np.asarray(d.loss_fn_inputs["weights"].tolist(), dtype=np.float64)
            m = w > 0; per.append((float(-lp[m].sum()), int(m.sum())))
            s += per[-1][0]; n += per[-1][1]
    return s, n, per

def retry(fn, what, waits=(60, 120, 240, 240, 240)):
    for i, w in enumerate((*waits, None)):
        try:
            return fn()
        except Exception as e:
            c = e.__cause__
            msg = f"{type(e).__name__}: {str(e)[:300]} cause={type(c).__name__ + ': ' + str(c)[:300] if c else None}"
            if w is None or not any(k in msg.lower() for k in ("capacity", "429", "503", "connection", "timeout", "unavailable")):
                raise
            print(f"retry {what} in {w}s: {msg}", flush=True); time.sleep(w)