"""N3 — shade map must drop the old floor's repeating tile wash (ghost
diagonal banding, room 4) while keeping real lighting: gradients and contact
shadows. Runs the REAL build_shade_map from app.py on a synthetic floor whose
components are known exactly, and compares against the pre-N3 pipeline.

Floor luminance = gradient * shadow * periodic tile wash * grout lines.
  - wash leakage  : correlation of the decoded shade with the wash component
                    -> must drop >= 60% vs the pre-N3 pipeline
  - shadow keep   : correlation with the shadow component
                    -> must stay >= 75% of the pre-N3 pipeline's
  - gradient keep : decoded left/right brightness ratio within 15% of truth
"""

import cv2
import numpy as np

# --- extract the real implementations from app.py ---------------------------
src = open("app.py").read()
ns = {"np": np, "cv2": cv2}
for fn in ["_adaptive_shade_range", "_encode_shade", "_dominant_period",
           "_suppress_periodic_shading", "build_shade_map"]:
    start = src.index(f"def {fn}")
    end = src.index("\ndef ", start + 10)
    # build_shade_map is followed by another def inside the same block scan
    exec(compile(src[start:end], "app.py", "exec"), ns)

build_shade_map = ns["build_shade_map"]


def build_shade_map_pre_n3(img_np, surface_mask):
    """The pre-N3 pipeline (median + Gaussian only), for the baseline."""
    mask = surface_mask.astype(np.uint8)
    luminance = (img_np[:, :, 0].astype(np.float32) * 0.299
                 + img_np[:, :, 1].astype(np.float32) * 0.587
                 + img_np[:, :, 2].astype(np.float32) * 0.114)
    h, w = mask.shape[:2]
    median_lum = float(np.median(luminance[mask > 0]))
    filled = luminance.copy()
    filled[mask == 0] = median_lum
    med_k = max(9, int(min(h, w) / 40)) | 1
    filled = cv2.medianBlur(np.clip(filled, 0, 255).astype(np.uint8), med_k).astype(np.float32)
    sigma = max(8.0, min(h, w) / 28.0)
    smooth = cv2.GaussianBlur(filled, (0, 0), sigmaX=sigma, sigmaY=sigma)
    relative = smooth / median_lum
    relative[mask == 0] = 1.0
    lo, hi = ns["_adaptive_shade_range"](relative, mask)
    return ns["_encode_shade"](relative, lo, hi), (lo, hi)


def decode(shade, rng):
    lo, hi = rng
    return lo + shade.astype(np.float32) / 255.0 * (hi - lo)


def masked_corr(a, b, m):
    av = a[m] - a[m].mean()
    bv = b[m] - b[m].mean()
    den = np.sqrt((av ** 2).sum() * (bv ** 2).sum()) + 1e-9
    return float((av * bv).sum() / den)


def main():
    H, W = 700, 900
    yy, xx = np.mgrid[0:H, 0:W].astype(np.float32)

    # known components
    gradient = 0.85 + 0.5 * (xx / W)                       # window on the right
    shadow = 1.0 - 0.35 * np.exp(-(((xx - 250) / 90) ** 2 + ((yy - 420) / 60) ** 2))
    period = 200
    wash = 1.0 + 0.16 * np.sign(np.sin(2 * np.pi * (xx + yy) / period)
                                * np.sin(2 * np.pi * (xx - yy) / period))
    wash = cv2.GaussianBlur(wash, (0, 0), 9)               # soft tile shading
    grout = np.where((np.mod(xx + yy, period) < 6) | (np.mod(xx - yy, period) < 6), 0.75, 1.0)

    lum = 150.0 * gradient * shadow * wash * grout
    img = np.repeat(np.clip(lum, 0, 255)[..., None], 3, axis=2).astype(np.uint8)
    mask = np.ones((H, W), np.uint8)
    mask[: H // 6] = 0                                     # a wall strip, exercises inpaint

    ok = True
    res = {}
    for name, fn in [("pre-N3", build_shade_map_pre_n3), ("N3", build_shade_map)]:
        shade, rng = fn(img, mask)
        rel = decode(shade, rng)
        m = mask.astype(bool)
        wash_c = masked_corr(rel, wash, m)
        shadow_c = masked_corr(rel, shadow, m)
        left = rel[m & (xx < W * 0.25)].mean()
        right = rel[m & (xx > W * 0.75)].mean()
        grad_ratio = left / right
        true_ratio = gradient[m & (xx < W * 0.25)].mean() / gradient[m & (xx > W * 0.75)].mean()
        res[name] = (wash_c, shadow_c, grad_ratio)
        print(f"[{name:6s}] wash-corr={wash_c:.3f}  shadow-corr={shadow_c:.3f}  "
              f"gradient L/R={grad_ratio:.3f} (truth {true_ratio:.3f})")

    wash_drop = 1 - abs(res["N3"][0]) / max(abs(res["pre-N3"][0]), 1e-6)
    shadow_keep = abs(res["N3"][1]) / max(abs(res["pre-N3"][1]), 1e-6)
    print(f"wash leakage drop = {wash_drop * 100:.0f}%  (need >= 60%)")
    print(f"shadow retention  = {shadow_keep * 100:.0f}%  (need >= 75%)")
    if wash_drop < 0.60:
        print("  !! periodic wash still leaking"); ok = False
    if shadow_keep < 0.75:
        print("  !! real shadow lost"); ok = False
    true_ratio = gradient[mask.astype(bool) & (xx < W * 0.25)].mean() / \
        gradient[mask.astype(bool) & (xx > W * 0.75)].mean()
    if abs(res["N3"][2] - true_ratio) > 0.15 * true_ratio:
        print("  !! lighting gradient distorted"); ok = False

    # period detector sanity on aperiodic field: pure shadow must NOT register
    aper = cv2.GaussianBlur((shadow * 40).astype(np.float32), (0, 0), 3)
    aper_hp = aper - cv2.GaussianBlur(aper, (0, 0), 50)
    if ns["_dominant_period"](aper_hp, 1) or ns["_dominant_period"](aper_hp, 0):
        print("  !! aperiodic shadow field misdetected as periodic"); ok = False

    print("\n" + ("ALL N3 CHECKS PASSED" if ok else "N3 CHECKS FAILED"))
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())