"""Create one fresh rank-32 LoRA, save it as the shared immutable init checkpoint, record a fixed-batch probe."""
import hashlib, json, time
import common, tinker
from tinker import types
from tinker_cookbook import renderers, checkpoint_utils
from tinker_cookbook.tokenizer_utils import get_tokenizer

MODEL, RENDERER = "Qwen/Qwen3.5-2B", "qwen3_5_disable_thinking"


def fixed_batch():
    tok = get_tokenizer(MODEL)
    r = renderers.get_renderer(RENDERER, tokenizer=tok)
    convo = [{"role": "user", "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Provide a numerical answer without units, written inside \\boxed{}."},
             {"role": "assistant", "content": "In May she sold 48/2 = 24 clips. Altogether 48 + 24 = 72. \\boxed{72}"}]
    mi, w = r.build_supervised_example(convo)
    toks = mi.to_ints()
    datum = types.Datum(model_input=types.ModelInput.from_ints(toks[:-1]),
                        loss_fn_inputs={"target_tokens": toks[1:], "weights": [float(x) for x in w[1:].tolist()]})
    h = hashlib.sha256(json.dumps({"tokens": toks[:-1], "targets": toks[1:], "weights": [float(x) for x in w[1:].tolist()]}).encode()).hexdigest()
    return datum, h


def probe(tc, datum):
    out = tc.forward([datum], "cross_entropy").result()
    lp = [float(x) for x in out.loss_fn_outputs[0]["logprobs"].tolist()]
    return lp, hashlib.sha256(json.dumps(lp).encode()).hexdigest()


if __name__ == "__main__":
    sc = common.service()
    md = {}
    checkpoint_utils.add_renderer_name_to_user_metadata(md, RENDERER)
    t0 = time.time()
    for i in range(20):
        try:
            tc = sc.create_lora_training_client(MODEL, rank=32, user_metadata=md)
            break
        except Exception as e:
            print("create failed", repr(e)[:300]); time.sleep(30)
    datum, bh = fixed_batch()
    lp, oh = probe(tc, datum)
    save = tc.save_state("rx-init-run_4125735d30").result()
    info = tc.get_info()
    rest = sc.create_rest_client()
    try:
        winfo = rest.get_weights_info_by_tinker_path(save.path).result()
        winfo = repr(winfo)
    except Exception as e:
        winfo = f"unavailable: {e!r}"[:300]
    res = {"init_checkpoint_uri": save.path, "save_response": repr(save), "training_client_info": repr(info),
           "weights_info": winfo, "fixed_batch_sha256": bh, "creator_output_sha256": oh,
           "creator_logprobs_first8": lp[:8], "seconds": time.time() - t0}
    json.dump(res, open("/tmp/rx-run/init.json", "w"), indent=1)
    print(json.dumps(res, indent=1))