"""Fit the canopy (tg_canopy.py) to the picture: the glass and its dark frame through the fitted camera against the region the
picture shows them in (traced polygon below) and against the frame's navy pixels.

python tg_fit_canopy.py   -> build/canopy.json, previews/compare/canopy_fit.png
"""
import json
import math
import numpy as np
from PIL import Image, ImageDraw
from tg_cam import load as load_cam, HERE
from tg_raster import raster
from tg_nm import nelder_mead
import tg_canopy as tc

CAM = load_cam()["params"]
ROOT = HERE.parent
SC = 0.5
# the glass with its frame as the picture shows it: the tip, the glass's top edge (under the brow from x 150), the frame's outer
# edge round the rear and along the lower flank, the glass's lower edge back to the tip
TARGET = [(14, 416), (20, 403), (30, 388), (40, 372), (60, 362), (80, 352), (100, 338), (120, 327), (140, 316), (160, 314),
          (180, 310), (200, 306), (220, 301), (240, 301), (260, 302), (272, 307), (280, 318), (292, 330), (303, 346),
          (308, 362), (310, 386), (304, 394), (299, 402), (287, 410), (280, 418), (269, 426), (255, 434), (240, 441),
          (225, 447), (210, 456), (195, 460), (180, 465), (165, 470), (150, 475), (135, 480), (120, 484), (100, 468),
          (80, 462), (60, 455), (40, 443), (30, 432), (20, 425), (14, 418)]
W, H = int(round(827 * SC)), int(round(803 * SC))
m = Image.new("L", (W, H), 0)
ImageDraw.Draw(m).polygon([(u * SC, v * SC) for u, v in TARGET], fill=1)
TGT = np.asarray(m).astype(bool)
cls = np.load(ROOT / "reference" / "classes.npy")
NAVY = np.asarray(Image.fromarray((cls == 7).astype(np.uint8) * 255).resize((W, H), Image.BILINEAR)) > 100
NAVY &= TGT

KEYS = ["y_rear", "w_max", "w_pow", "crest", "bot", "mid_frac", "y_jaw", "sill_rear", "sill_mid"]


def to_vec(p):
    v = []
    for k in KEYS:
        v += list(p[k]) if isinstance(p[k], list) else [p[k]]
    return np.array(v, float)


def to_dict(v, base):
    p = dict(base)
    i = 0
    for k in KEYS:
        if isinstance(base[k], list):
            p[k] = [float(x) for x in v[i:i + len(base[k])]]
            i += len(base[k])
        else:
            p[k] = float(v[i])
            i += 1
    return p


def render(p):
    V, F, path, info = tc.canopy(p, n_st=22, n_around=32)
    fv, ff = tc.frame_band(path, p)
    tris, ids = [], []
    for verts, faces, k in ((V, F, 1), (fv, ff, 2)):
        for f in faces:
            for i in range(1, len(f) - 1):
                tris.append(verts[[f[0], f[i], f[i + 1]]])
                ids.append(k)
    img, _ = raster(CAM, np.array(tris), np.array(ids), SC)
    return img


def loss(v, base):
    p = to_dict(v, base)
    if not (0.35 < p["w_max"] < 0.95 and 0.4 < p["w_pow"] < 2.5 and 0.2 < p["mid_frac"] < 0.8
            and p["y_nose"] + 0.3 < p["y_jaw"] < p["y_rear"] - 0.3 and -1.2 < p["y_rear"] < 0.2):
        return 2.0
    img = render(p)
    mm = img > 0
    iou = (mm & TGT).sum() / max(1, (mm | TGT).sum())
    fr = img == 2
    iou_f = (fr & NAVY).sum() / max(1, (fr | NAVY).sum())
    # keep the crest rising and above the bottom
    c = [p["z_nose"]] + p["crest"]
    pen = sum(max(0, c[i] - c[i + 1] - 0.02) for i in range(len(c) - 1))
    return (1 - iou) + 0.35 * (1 - iou_f) + pen


def main():
    base = tc.load()
    v0 = to_vec(base)
    step = np.full(len(v0), 0.05)
    best, val = None, None
    for rnd in range(3):
        v, f = nelder_mead(lambda x: loss(x, base), v0, step, iters=350,
                           log=lambda it, b: print("  it", it, round(b, 4), flush=True))
        print("round", rnd, "loss", round(f, 4), flush=True)
        if val is None or f < val:
            best, val = v, f
        v0, step = best, step * 0.5
    p = to_dict(best, base)
    keep = {k: p[k] for k in KEYS}
    (HERE / "canopy.json").write_text(json.dumps(keep, indent=1))
    img = render(p)
    pic = Image.open(ROOT / "reference" / "user_picture.jpg").convert("RGB").resize((W, H))
    a = np.asarray(pic).copy()
    a[(img > 0) & ~TGT] = (255, 0, 0)
    a[~(img > 0) & TGT] = (0, 0, 255)
    a[img == 2] = (a[img == 2] * 0.5 + np.array([255, 255, 0]) * 0.5).astype(np.uint8)
    Image.fromarray(a).resize((827, 803)).save(ROOT / "previews" / "compare" / "canopy_fit.png")
    print(json.dumps(keep))


main()