"""Render triangles through the picture's camera as flat colour classes and compare with the picture.

compare(tris, cls, tag) -> previews/compare/<tag>_classes.png (picture classes | model classes),
                           previews/compare/<tag>_over.png (the model's class boundaries over the picture)
Class ids as tg_classes.CLASSES; 8 = glass, 9 = other (steel and small fittings).
"""
import numpy as np
from PIL import Image
from tg_cam import load, HERE
from tg_raster import raster, edges_of
from tg_classes import COLOURS

ROOT = HERE.parent
CMP = ROOT / "previews" / "compare"
PAL = np.array(COLOURS + [(150, 225, 240), (250, 250, 120)], np.uint8)


def compare(tris, cls, tag, cam=None, scale=1.0):
    cam = cam or load()["params"]
    img, dep = raster(cam, np.asarray(tris, float), np.asarray(cls, np.int32), scale)
    pic = Image.open(ROOT / "reference" / "user_picture.jpg").convert("RGB")
    pcls = np.load(ROOT / "reference" / "classes.npy")
    model = PAL[np.where(img < 0, 0, img)]
    model[img < 0] = (235, 235, 235)
    both = np.concatenate([PAL[pcls], np.full((pcls.shape[0], 6, 3), 255, np.uint8), model], 1)
    Image.fromarray(both).save(CMP / f"{tag}_classes.png")
    over = np.asarray(Image.blend(pic, Image.new("RGB", pic.size, (255, 255, 255)), 0.25)).copy()
    e = edges_of(img)
    over[e] = (255, 0, 255)
    sil = edges_of((img >= 0).astype(np.int32))
    over[sil] = (0, 200, 0)
    Image.fromarray(over).save(CMP / f"{tag}_over.png")
    fg_m, fg_p = img >= 0, pcls > 0
    iou = (fg_m & fg_p).sum() / max(1, (fg_m | fg_p).sum())
    out = {"fg_iou": round(float(iou), 3)}
    for k, n in ((1, "orange"), (2, "red")):
        a, b = img == k, pcls == k
        out[n + "_iou"] = round(float((a & b).sum() / max(1, (a | b).sum())), 3)
    print(tag, out)
    return img, out


MAKO_CLASS = {"Hull paint": 1, "Jaw paint": 2, "Fin grey": 3, "Side body": 3, "Recess grey": 3, "Lower hull": 4,
              "Graphite": 4, "Rubber": 4, "Black gloss": 5, "Glass": 8, "Glass upper": 8, "Steel": 9,
              "Yellow marking": 9, "Lamp light": 9}

if __name__ == "__main__":
    d = np.load(HERE / "checks" / "mako_tris.npz")
    mats = [m.replace("MAKO | ", "") for m in d["mats"]]
    cls = np.array([MAKO_CLASS.get(m, 4) for m in mats])
    keep = np.array([not m.startswith("Cabin") and m not in ("Seat", "Screen") for m in mats])
    compare(d["tris"][keep], cls[keep], "mako")