"""Record a reference encoding + reconstruction of the bundled images, for the stability tests.

What is being pinned is the PICTURE, not the codes. Codes are an argmin over 4096 near-tied
codebook entries; fp16 flips about 0.5% of them against fp32. Measured, that costs 0.47/255 mean
absolute pixel difference -- against a reconstruction error of 10.25/255 -- and the differences sit
on high-frequency edges, where a code flip moves an edge by a fraction of a pixel and a single
pixel can swing 159/255 without anything semantic changing. So the reference stores both, and the
tests judge on mean pixel difference.

JSON plus PNG, no binary blobs: a reference is something a human should be able to open and diff
when a new device disagrees.

  python -m satlink.tools.make_reference
"""
import argparse
import glob
import hashlib
import json
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from satlink.frame import VQFrame                 # noqa: E402
from satlink.protocol import LinkProfile          # noqa: E402
from satlink.sat.encoder import SatEncoder        # noqa: E402



def _pixel_grid(model, prof):
    f = model.patch_size * model.merge
    return (max(f, round(prof.height / f) * f) // model.patch_size,
            max(f, round(prof.width / f) * f) // model.patch_size)


def _teacher_forced_logprobs(model, frame, ids, pos, answer, device):
    """-> the logprob the model assigns to each answer token, given the prompt and all earlier
    answer tokens. One forward pass, deterministic, no decoding."""
    import torch
    full = torch.cat([ids, torch.tensor(answer, device=device)])
    nxt = int(pos.max()) + 1
    tail = torch.arange(nxt, nxt + len(answer), device=device).expand(3, -1)
    fpos = torch.cat([pos, tail], dim=1)
    with torch.no_grad(), model.reference_attention():
        h = model(full, fpos, frame=frame)
        logits = model.compute_logits(h[len(ids) - 1:-1]).float()
    lp = torch.log_softmax(logits, -1)
    return lp[torch.arange(len(answer)), torch.tensor(answer, device=device)].tolist()


def _record_answers(ref, a):
    """Record, per image: the generated answer, and the logprob the model assigns to each of its
    tokens under teacher forcing.

    The test uses BOTH. Teacher forcing is the stable measurement -- one forward pass, no decoding,
    a continuous number that moves smoothly with numerical drift. Free generation is the one that
    tells you what a user would actually see, but it is brittle by nature: greedy decoding turns
    any near-tie into a different trajectory, so it is reported as a common-prefix length rather
    than an equality check.
    """
    import torch
    from tokenizers import Tokenizer
    from satlink.frame import VQFrame
    from satlink.ground.infer import greedy
    from satlink.ground.modeling_vqvlm import mrope_positions

    if a.latent_ground:
        from satlink.ground.modeling_latentvlm import LatentGlm4vForConditionalGeneration as M
        model = M.from_pretrained(a.model, a.latent_ground, lora=a.lora, device=a.device)
        tail, path = "<|assistant|>", "latent"
    else:
        from satlink.ground.modeling_vqvlm import VQGlm4vForConditionalGeneration as M
        model = M.from_pretrained(a.model, a.decoder, LinkProfile(512, 512), device=a.device)
        tail, path = "<|assistant|>\n", "pixel"
    tok = Tokenizer.from_file(os.path.join(a.model, "tokenizer.json"))
    IMG = model.image_token_id
    ref["answers"] = {"path": path, "question": a.question, "items": {}}

    for name, meta in sorted(ref["images"].items()):
        prof = LinkProfile(meta["width"], meta["height"], n_tok=a.n_tok)
        frame = VQFrame(torch.tensor(meta["codes"], dtype=torch.long), prof).to(a.device)
        n_img = model.n_image_tokens(prof) if a.latent_ground else model.n_image_tokens()
        gh, gw = model.image_grid(prof) if a.latent_ground else _pixel_grid(model, prof)
        ids = torch.tensor(
            tok.encode("[gMASK]<sop><|user|>\n<|begin_of_image|>", add_special_tokens=False).ids
            + [IMG] * n_img
            + tok.encode(f"<|end_of_image|>{a.question}{tail}", add_special_tokens=False).ids,
            device=a.device)
        pos = mrope_positions(ids, IMG, gh, gw, model.merge).to(a.device)
        with model.reference_attention():
            ans = greedy(model, ids, pos, frame, max_new=a.answer_tokens)
        lp = _teacher_forced_logprobs(model, frame, ids, pos, ans, a.device)
        ref["answers"]["items"][name] = {
            "tokens": ans, "logprobs": [round(x, 4) for x in lp],
            "text": tok.decode(ans, skip_special_tokens=False)}
        print(f"  answer {name:24} {len(ans)} tok, mean logprob {sum(lp)/len(lp):.4f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--assets", default="satlink/assets")
    ap.add_argument("--sat", default="satlink/weights/sat_encoder.pt")
    ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt")
    ap.add_argument("--n-tok", type=int, default=48)
    ap.add_argument("--model", default=None,
                    help="also record a reference ANSWER and its teacher-forced logprobs")
    ap.add_argument("--latent-ground", default=None)
    ap.add_argument("--lora", default=None)
    ap.add_argument("--question", default="where is the largest building in the image")
    ap.add_argument("--answer-tokens", type=int, default=96)
    ap.add_argument("--device", default="cuda")
    a = ap.parse_args()

    from PIL import Image
    import torchvision.transforms.functional as TF
    from satlink.ground.decoder import GroundDecoder
    gnd = GroundDecoder(a.decoder).to(a.device)

    ref = {"n_tok": a.n_tok, "dtype": "float32",
           "device": torch.cuda.get_device_name(0) if a.device.startswith("cuda") else "cpu",
           "note": "codes recorded in fp32; tests compare RECONSTRUCTED PIXELS, not codes",
           "images": {}}
    os.makedirs(os.path.join(a.assets, "reference"), exist_ok=True)
    for path in sorted(glob.glob(os.path.join(a.assets, "*.png"))):
        img = Image.open(path).convert("RGB")
        prof = LinkProfile(*img.size, n_tok=a.n_tok)
        sat = SatEncoder(a.sat, prof, device=a.device, dtype=torch.float32)
        codes = sat.encode(img)
        with torch.no_grad():
            rec = gnd(VQFrame(codes.tokens, prof).to(a.device)).cpu()
        name = os.path.basename(path)
        rec_path = os.path.join(a.assets, "reference", name)
        TF.to_pil_image(rec).save(rec_path, "PNG", optimize=True)
        ref["images"][name] = {
            "sha256": hashlib.sha256(open(path, "rb").read()).hexdigest()[:16],
            "width": img.size[0], "height": img.size[1], "grid": list(prof.grid),
            "tile_origins": [prof.tile_origin(i) for i in range(prof.n_tiles)],
            "recon": f"reference/{name}",
            "codes": codes.tolist(),
        }
        print(f"  {name:26} {img.size[0]}x{img.size[1]:<5} {prof.n_tiles} crops -> "
              f"{prof.frame_bytes} B, recon {os.path.getsize(rec_path)/1e3:.0f} KB")
    if a.model:
        _record_answers(ref, a)

    out = os.path.join(a.assets, "reference.json")
    json.dump(ref, open(out, "w"), indent=1)
    print(f"-> {out}  ({os.path.getsize(out)/1e6:.2f} MB, recorded on {ref['device']})")
    return 0


if __name__ == "__main__":
    sys.exit(main())