"""Compare the old centre-crop with the letterbox on real non-square photos.

There are no labels for these photos here, so the reference is the model's own
prediction, and the numbers are *differences*, read against yardsticks:
  crop vs crop-run-twice   the ONNX export samples a little input noise per call
  crop vs mirrored run     a benign test-time variation (the mirror average)
  crop vs letterbox        the thing being tested, on the shared central region
  edge band vs inner       does reflected padding hurt next to the padding?
Depth is also compared after a least-squares scale+shift (its scale is relative).

  docker compose run --rm train bash -c "pip install -q onnxruntime huggingface_hub && python space/test_letterbox.py"
"""
import glob
import os
import sys

import numpy as np
import onnxruntime as ort
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from framing import crop_to_box, letterbox, resize_maps  # noqa: E402

SIZE = 448
OUT = os.path.join(os.path.dirname(HERE), "_lb_out")
os.makedirs(OUT, exist_ok=True)

so = ort.SessionOptions()
so.log_severity_level = 3
so.intra_op_num_threads = 4
SESS = ort.InferenceSession(
    hf_hub_download("singam96/ShadeNet-5-35M", "onnx/shadenet5_mix_int4_int8.onnx"),
    so, providers=["CPUExecutionProvider"])


def tensor(img):
    a = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(a.transpose(2, 0, 1)[None] * 2 - 1)


def run(x):
    return torch.from_numpy(SESS.run(None, {"input_rgb": np.ascontiguousarray(x.numpy())})[0])


def centre_crop(img):                     # the old behaviour
    w, h = img.size
    s = max(SIZE / w, SIZE / h)
    nw, nh = int(round(w * s)), int(round(h * s))
    im = img.convert("RGB").resize((nw, nh), Image.BICUBIC)
    l, t = (nw - SIZE) // 2, (nh - SIZE) // 2
    return im.crop((l, t, l + SIZE, t + SIZE))


def mirrored(x):
    m = run(x.flip(-1)).flip(-1).clone()
    m[:, 4] = -m[:, 4]
    return m


def ssi(a, b):
    """Depth a aligned to b by least-squares scale+shift."""
    A = torch.stack([a.flatten(), torch.ones(a.numel())], 1)
    sol = torch.linalg.lstsq(A, b.flatten()[:, None]).solution
    return (A @ sol).reshape(a.shape)


def diffs(a, b):
    return (float((a[:, 0:3] - b[:, 0:3]).abs().mean()),
            float((a[:, 3:4] - b[:, 3:4]).abs().mean()),
            float((ssi(a[:, 3:4], b[:, 3:4]) - b[:, 3:4]).abs().mean()),
            float((a[:, 4:7] - b[:, 4:7]).abs().mean()))


def tile(x01, ts):
    if x01.shape[1] == 1:
        x01 = x01.expand(-1, 3, -1, -1)
    return Image.fromarray((x01[0].permute(1, 2, 0).clamp(0, 1).numpy() * 255).round().astype("uint8")).resize(ts, Image.BICUBIC)


rows = []
for f in sorted(glob.glob(os.path.join(os.path.dirname(HERE), "dist", "hf_bundle", "assets", "examples", "*.jpg"))):
    img = Image.open(f).convert("RGB")
    w, h = img.size
    xc = tensor(centre_crop(img))
    mc, mc2, mm = run(xc), run(xc), mirrored(xc)
    sq, box = letterbox(img, SIZE)
    ml = crop_to_box(run(tensor(sq)), box)                    # [1,7,nh,nw] real photo only
    left, top, nw, nh = box
    # shared region: what the crop saw, in the letterbox photo's coordinates
    if w >= h:
        c0, c1 = int(round(nw * (w - h) / (2 * w))), int(round(nw * (w + h) / (2 * w)))
        shared = ml[:, :, :, c0:c1]
        edge = torch.zeros(nh, c1 - c0, dtype=torch.bool)
        k = max(1, int(0.125 * nh)); edge[:k] = True; edge[-k:] = True
    else:
        r0, r1 = int(round(nh * (h - w) / (2 * h))), int(round(nh * (h + w) / (2 * h)))
        shared = ml[:, :, r0:r1, :]
        edge = torch.zeros(r1 - r0, nw, dtype=torch.bool)
        k = max(1, int(0.125 * nw)); edge[:, :k] = True; edge[:, -k:] = True
    shared = F.interpolate(shared, size=(SIZE, SIZE), mode="bicubic", align_corners=False).clamp(-1, 1)
    edge_up = F.interpolate(edge[None, None].float(), size=(SIZE, SIZE), mode="nearest")[0, 0] > 0.5
    d_noise, d_mirror, d_lb = diffs(mc, mc2), diffs(mc, mm), diffs(shared, mc)
    inner = (~edge_up)[None, None]
    e_in = float(((shared - mc).abs() * inner).sum() / (inner.sum() * 7))
    e_edge = float(((shared - mc).abs() * edge_up[None, None]).sum() / (edge_up.sum() * 7))
    rows.append((os.path.basename(f)[:10], f"{w}x{h}", d_noise, d_mirror, d_lb, e_edge, e_in))
    # picture: letterboxed input + its maps (padding cropped off) over the crop's maps
    ts = (nw, nh)
    top_row = [Image.fromarray(np.asarray(sq)[top:top + nh, left:left + nw])] + [tile((ml[:, a:b] + 1) / 2, ts) for a, b in ((0, 3), (3, 4), (4, 7))]
    s = 360 / max(nw, nh)
    cw, ch = int(nw * s), int(nh * s)
    canvas = Image.new("RGB", (4 * cw, ch + 360), (24, 24, 24))
    for i, t in enumerate(top_row):
        canvas.paste(t.resize((cw, ch)), (i * cw, 0))
    canvas.paste(centre_crop(img).resize((360, 360)), (0, ch))
    for i, (a, b) in enumerate(((0, 3), (3, 4), (4, 7))):
        canvas.paste(tile((mc[:, a:b] + 1) / 2, (360, 360)), ((i + 1) * 360, ch))
    canvas.save(os.path.join(OUT, os.path.basename(f).replace(".jpg", ".png")))

print(f"{'image':11s} {'size':8s} | mean |diff| over 7 maps: alb/depth/depth-aligned/normal")
hdr = ["noise (run twice)", "crop vs mirrored", "crop vs letterbox"]
for r in rows:
    print(f"{r[0]:11s} {r[1]:8s} | " + " | ".join(f"{n}: " + "/".join(f"{v:.3f}" for v in d) for n, d in zip(hdr, r[2:5])))
mean = lambda i: np.mean([np.array(r[i]) for r in rows], axis=0)
print("\nMEAN over photos (alb/depth/depth-aligned/normal)")
for n, i in zip(hdr, (2, 3, 4)):
    print(f"  {n:20s}", "/".join(f"{v:.4f}" for v in mean(i)))
print(f"\nletterbox vs crop, mean over all 7 channels: edge band {np.mean([r[5] for r in rows]):.4f}  inner {np.mean([r[6] for r in rows]):.4f}")
print("pictures in", OUT)