"""Check the Space's copied averaging code against the repository's original.

Run from the repo root inside the project container:
  docker compose run --rm train bash -c "pip install -q onnxruntime gradio huggingface_hub && python space/test_space.py"
"""
import os
import sys
import time

import torch
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(HERE)
sys.path.insert(0, HERE)
os.chdir(HERE)

import app  # noqa: E402
import ensemble  # noqa: E402  (downloads config.json + the MIX ONNX from the Hub)

sys.path.insert(0, REPO)
from detail_diffuser.inference.core import tta_average as ref_tta, unit_normals as ref_unit  # noqa: E402

img = Image.open(os.path.join(HERE, "examples", "3185371756_ff4e9fa8a6.jpg"))
from framing import letterbox  # noqa: E402
sq, box = letterbox(img, app.SIZE)
x = app.to_tensor(sq)
assert x.shape == (1, 3, app.SIZE, app.SIZE) and x.min() >= -1 and x.max() <= 1
assert box[2] == app.SIZE or box[3] == app.SIZE, box   # the long side fills the square

t = time.time()
fast = app.refine(x)
print(f"fast: {time.time() - t:.1f}s, maps {tuple(fast.shape)}")
assert fast.shape == (1, 7, app.SIZE, app.SIZE) and torch.isfinite(fast).all()

# The exported graph samples its input-perturbation noise on every call, so
# two runs of the real model differ slightly (mean ~1e-3). The averaging code
# is therefore compared with a deterministic stand-in model, which must match
# the repository's implementation exactly.
def stand_in(c):
    return torch.cat([c * 0.7, c.mean(1, keepdim=True) * 0.5, c.flip(-1) * 0.3], 1)


mine = app.unit_normals(ensemble.tta_average(stand_in, x, flip=True, rot_deg=app.ROT_DEG))
ref = ref_unit(ref_tta(lambda c, nz=None: stand_in(c), x, flip=True, rot_deg=app.ROT_DEG))
d = float((mine - ref).abs().max())
print("averaging code, max |space - repo| =", d)
assert d == 0.0, d

e1 = app.unit_normals(ensemble.tta_average(app.refine, x, flip=True, rot_deg=app.ROT_DEG))
assert e1.shape == (1, 7, app.SIZE, app.SIZE) and torch.isfinite(e1).all()
n = e1[:, 4:7].norm(dim=1)
assert float((n - 1).abs().max()) < 1e-4, "normals are not unit length"
print("real model, 6 passes ok, mean |maps| =", float(e1.abs().mean()))

# median_of against NumPy's nanmedian on real stand-in passes (corners are covered by fewer passes)
import numpy as np  # noqa: E402

passes = app.tta_passes(lambda c: stand_in(c) + 0.05 * torch.randn(1, 7, app.SIZE, app.SIZE),
                        x, flip=True, rot_deg=app.ROT_DEG)
assert len(passes) == 6 and any(float(v.min()) < 0.5 for _, v in passes), "rotated passes should have uncovered corners"
P = np.stack([p.numpy() for p, _ in passes]).copy()
V = np.stack([v.numpy() for _, v in passes]) > 0.5
P[~np.broadcast_to(V, P.shape)] = np.nan
want = np.nanmedian(P, axis=0)
got = app.median_of(passes).numpy()
dm = float(np.abs(got - want).max())
print("median_of vs numpy nanmedian, max diff =", dm)
assert dm < 1e-5, dm
assert float((app.median_of(passes) - app.mean_of(passes)).abs().mean()) > 1e-4, "median should differ from mean on noisy passes"
one = [passes[0]]
assert torch.equal(app.median_of(one), app.mean_of(one)), "one pass: median == mean"

out = app.predict(img, app.MODE_SIX, False)
assert len(out) == 6 and out[0].size == out[1].size == img.size, (out[0].size, img.size)  # whole photo, own shape
assert out[4] is None and "(mean)" in out[5]
print("predict ok:", out[5])
out = app.predict(img, app.MODE_SIX, True)
assert out[4] is not None and out[4].size == img.size and "(median)" in out[5], out[5]
print("median ok:", out[5])
out = app.predict(img, app.MODE_SINGLE, True)        # median outside the six-pass mode: single pass, note says so
assert out[4] is None and "1 forward pass in" in out[5] and "six-pass mode" in out[5], out[5]
out = app.predict(img, app.MODE_REFINE, False)       # single pass + the small refiner
assert out[4] is None and "refiner" in out[5] and out[1].size == img.size, out[5]
print("refiner mode ok:", out[5])
# odd shapes: a panorama (padding longer than the photo), tall, square, tiny, big
for size in [(2000, 300), (300, 2000), (448, 448), (64, 48), (1500, 1100)]:
    im = Image.effect_noise(size, 40).convert("RGB")
    out = app.predict(im, app.MODE_REFINE, False)
    ow, oh = app.display_size(*size)
    assert out[0].size == out[1].size == out[2].size == out[3].size == (ow, oh), (size, out[0].size, (ow, oh))
print("odd shapes ok")
print("space test ok")