"""GROUND side, tokenizer only: downlink codes in, a picture out. No VLM.

The counterpart to sat_encode.py. Between them they are the whole link, and each runs alone --
this one never imports satellite code, and a codes file is all it needs.

    python -m satlink.examples.sat_encode  --image scene.png --out codes.json
    python -m satlink.examples.ground_decode --codes codes.json --out recon.png

Useful on its own for the thing that actually matters when a new device joins the constellation:
whether its codes still make the same picture. Pass `--compare` with a reference reconstruction and
it reports the pixel difference.

    python -m satlink.examples.ground_decode --codes codes.json --out recon.png \
        --compare satlink/assets/reference/grid3x3_600x600.png

Judge that on the MEAN difference. Codes are an argmin over 4096 near-tied codebook entries, so
fp16 flips about 0.5% of them against fp32; measured, that is 0.47/255 mean against a
reconstruction error of 10.25/255, and it lands on high-frequency edges where a single pixel can
swing 159/255 without anything semantic changing.
"""
import argparse
import json
import os
import sys
import time

import torch

from ..frame import VQFrame
from ..protocol import LinkProfile, unpack_frame


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--codes", required=True, help="codes.json (or packed bytes) from sat_encode")
    ap.add_argument("--out", required=True, help="where to write the reconstructed PNG")
    ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt")
    ap.add_argument("--width", type=int, default=0, help="agreed frame width (packed input only)")
    ap.add_argument("--height", type=int, default=0)
    ap.add_argument("--n-tok", type=int, default=48)
    ap.add_argument("--compare", default=None, help="a reference image to report pixel stats against")
    ap.add_argument("--device", default="cuda")
    a = ap.parse_args()

    import torchvision.transforms.functional as TF
    from ..ground.decoder import GroundDecoder

    # JSON is the readable form; packed bytes are what a radio carries. Either way the link layer
    # ends HERE -- the decoder takes a VQFrame.
    if a.codes.endswith(".json"):
        j = json.load(open(a.codes))
        profile = LinkProfile(j["width"], j["height"], n_tok=j["n_tok"])
        codes = j["codes"]
        src = f"{a.codes} ({j.get('dtype', '?')}, from {j.get('image', '?')})"
    else:
        if not (a.width and a.height):
            ap.error("--width/--height are required for packed input; they are the agreed profile")
        profile = LinkProfile(a.width, a.height, n_tok=a.n_tok)
        codes = unpack_frame(profile, open(a.codes, "rb").read())
        src = a.codes
    frame = VQFrame(torch.tensor(codes, dtype=torch.long), profile)
    tx, ty = profile.grid
    print(f"downlink   {src}")
    print(f"           {profile.frame_bytes} bytes, {tx}x{ty} crops x {profile.n_tok} tokens "
          f"-> {profile.width}x{profile.height}")

    gnd = GroundDecoder(a.decoder).to(a.device)
    t0 = time.time()
    with torch.no_grad():
        rec = gnd(frame.to(a.device)).cpu()
    dt = time.time() - t0
    print(f"decoder    {sum(p.numel() for p in gnd.parameters())/1e6:.1f}M params, {1000*dt:.0f} ms")

    os.makedirs(os.path.dirname(os.path.abspath(a.out)) or ".", exist_ok=True)
    TF.to_pil_image(rec).save(a.out, "PNG", optimize=True)
    print(f"output     {a.out}  ({os.path.getsize(a.out)/1e3:.0f} KB)")

    if a.compare:
        import numpy as np
        from PIL import Image
        want = torch.from_numpy(
            np.asarray(Image.open(a.compare).convert("L")).copy()).float() / 255.0
        if want.shape != rec[0].shape:
            print(f"compare    SKIPPED: {a.compare} is {tuple(want.shape)}, "
                  f"this reconstruction is {tuple(rec[0].shape)}")
            return 0
        d = (rec[0] - want).abs()
        print(f"compare    vs {a.compare}")
        print(f"           mean |Δ| {255*float(d.mean()):.3f}/255   max {255*float(d.max()):.1f}   "
              f">32/255 on {100*float((d > 32/255).float().mean()):.3f}% of pixels")
        print(f"           judge on the mean: reconstruction error itself is ~10.25/255, and "
              f"fp16-vs-fp32 code drift costs ~0.47")
    return 0


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