# End-to-End-Sim: App-Layouts (wie OsState.defaultLayout) -> App-Hint-Logik
# -> ONNX -> Glass-Frame. Verifiziert, was die App auf dem Handy zeigen wird.
import sys
import numpy as np
import onnxruntime as ort
from PIL import Image

W, H = 512.0, 1088.0
T = {"panel": 0.15, "card": 0.25, "circle": 0.35, "button": 0.45, "pill": 0.60, "dot": 0.85}


def app_layout(screen):
    els = []
    def el(kind, x0, y0, x1, y1, r):
        els.append(dict(kind=kind, x0=x0, y0=y0, x1=x1, y1=y1, r=r, t=T[kind]))
    el("dot", 28, 20, 52, 44, 10); el("pill", 240, 18, 336, 46, 12); el("dot", 448, 20, 472, 44, 10)
    el("pill", 124, 1012, 236, 1068, 26); el("pill", 196, 1024, 316, 1056, 14); el("pill", 352, 1012, 464, 1068, 26)
    if screen == "home":
        el("panel", 28, 68, 484, 240, 36)
        el("card", 28, 264, 300, 436, 28); el("card", 320, 264, 484, 436, 28)
        el("panel", 28, 456, 484, 576, 28)
        for i in range(4):
            x = 52 + i * 112
            el("pill", x, 880, x + 88, 968, 24)
    elif screen == "lock":
        el("circle", 116, 140, 396, 420, 140)
        el("panel", 96, 600, 416, 712, 40)
        el("button", 156, 780, 356, 868, 32)
    else:  # settings
        for i in range(4):
            el("card", 28, 72 + i * 124, 484, 72 + i * 124 + 100, 24)
        el("panel", 28, 584, 484, 712, 28)
    return els


def sdf_rr(fx, fy, e):
    cx, cy = (e["x0"] + e["x1"]) / 2, (e["y0"] + e["y1"]) / 2
    hx = (e["x1"] - e["x0"]) / 2 - e["r"]; hy = (e["y1"] - e["y0"]) / 2 - e["r"]
    qx = abs(fx - cx) - hx; qy = abs(fy - cy) - hy
    ax = np.maximum(qx, 0); ay = np.maximum(qy, 0)
    return np.sqrt(ax * ax + ay * ay) + np.minimum(np.maximum(qx, qy), 0) - e["r"]


def app_hint(els, ripples, hw=256, hh=544):
    hint = np.zeros((hh, hw, 3), np.float32)
    Y, X = np.mgrid[0:hh, 0:hw].astype(np.float32)
    fx = X * (W / hw) + 1.0; fy = Y * (H / hh) + 1.0   # Pixelzentren ~ Modellkoords
    for (rx, ry, age, s) in ripples:
        cx, cy = rx * W, ry * H
        dist = np.sqrt((X * (W / hw) - cx) ** 2 + (Y * (H / hh) - cy) ** 2)
        r0 = age * 300.0
        fade = (1.0 - age) * s
        ring = np.exp(-((dist - r0) ** 2) / 128.0) * fade
        glow = np.exp(-(dist ** 2) / 4608.0) * fade * 0.8
        hint[..., 1] += np.clip(ring + glow, 0, 1)
    for e in els:
        if e["kind"] == "circle":
            cx, cy = (e["x0"] + e["x1"]) / 2, (e["y0"] + e["y1"]) / 2
            sdf = np.sqrt((fx - cx) ** 2 + (fy - cy) ** 2) - (e["x1"] - e["x0"]) / 2
        else:
            sdf = sdf_rr(fx, fy, e)
        m_in = sdf <= 0
        border = np.abs(sdf) < 2.2
        t = e["t"]
        hint[..., 0] = np.where(border, np.maximum(hint[..., 0], 0.95), hint[..., 0])
        hint[..., 1] = np.where(border, np.maximum(hint[..., 1], 0.5), hint[..., 1])
        fill = max(0.28, t) if t < 0.45 else t
        hint[..., 1] = np.where(m_in, np.maximum(hint[..., 1], fill), hint[..., 1])
        hint[..., 2] = np.where(m_in, np.maximum(hint[..., 2], t), hint[..., 2])
    return hint


def stvec(screen_idx, ripples, phase=0.3, hue=210):
    v = np.zeros(32, np.float32)
    v[screen_idx] = 1.0
    for k, (rx, ry, rr_, s_) in enumerate(ripples[:2]):
        v[10 + k * 4] = rx; v[11 + k * 4] = ry; v[12 + k * 4] = rr_; v[13 + k * 4] = s_
    v[18] = np.sin(2 * np.pi * phase); v[19] = np.cos(2 * np.pi * phase)
    v[20] = 0.5; v[21] = np.sin(2 * np.pi * hue / 360); v[22] = np.cos(2 * np.pi * hue / 360)
    v[23] = 0.8; v[24] = 0.6; v[25] = 1.0; v[29] = 1.0; v[31] = 0.2
    return v


sess = ort.InferenceSession("/home/z/my-project/eos-v5/models/os_glass_v5.onnx",
                            providers=["CPUExecutionProvider"])
rows = []
for si, (name, idx) in enumerate([("lock", 0), ("home", 1), ("settings", 6)]):
    ripples = [(0.5, 0.45, 0.45, 1.0)] if name != "settings" else [(0.3, 0.75, 0.25, 0.9)]
    els = app_layout(name)
    hint = app_hint(els, ripples)                    # 544x256x3 (App-Res)
    vec = stvec(idx, ripples)
    out = sess.run(None, {"hint": np.transpose(hint, (2, 0, 1))[None],
                          "state": vec[None]})[0][0]
    frame = np.transpose(out, (1, 2, 0))
    hint_up = np.repeat(np.repeat(hint, 2, 0), 2, 1)
    row = np.concatenate([hint_up, frame], axis=1)
    rows.append(row)
    Image.fromarray((np.clip(frame, 0, 1) * 255).astype(np.uint8)).save(
        "/home/z/my-project/eos-v5/previews/app_%s.png" % name)
grid = (np.clip(np.concatenate(rows, 0), 0, 1) * 255).astype(np.uint8)
Image.fromarray(grid).save("/home/z/my-project/eos-v5/previews/app_sim_grid.png")
print("app sim ok -> previews/app_sim_grid.png")