#!/usr/bin/env python3
"""Merjenje kontrasta med svetlimi in temnimi kvadrati v 2x2 mrežah.

Slike tekočekristalnih struktur med prekrižanima polarizatorjema vsebujejo več
2x2 mrež. V vsaki mreži imata diagonalna para kvadratov paralelno oz. zvito
(twist) ureditev, zato sta dva kvadrata svetla in dva temna.

Uporaba:
  ./analyze.py                 # izbor manjkajočih mrež, nato izračun in grafi
  ./analyze.py --reselect      # ponovni izbor na vseh (navedenih) slikah
  ./analyze.py --compute-only  # samo izračun iz rois.json

Izbor: z miško povleci pravokotnik okoli cele 2x2 mreže, osemkrat na sliko, po
naraščajočem indeksu gostote (indeks 0..7, pitch = (indeks+1) * 100 nm).
Tipke: m = zadnji mreži preklopi oznako "manjka", u = razveljavi zadnjo,
enter = potrdi (ko jih je 8), esc = opusti sliko.

Mreže z "missing": true v rois.json niso izmerjene: ne pridejo v results.json
in ne na grafe, na pregledni sliki so sivo označene z napisom "manjka".

Upoštevana je EXIF orientacija, zato so koordinate v rois.json v isti
usmeritvi, kot jo pokaže pregledovalnik slik.

Izhod: rois.json, results.json, <slika>_overlay.png in štirje grafi.
"""

import argparse
import json
import re
import sys
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle
from matplotlib.widgets import RectangleSelector
from PIL import Image, ImageOps

N_GRIDS = 8  # število 2x2 mrež na sliko
PITCH_STEP_NM = 100  # pitch = (indeks + 1) * PITCH_STEP_NM
QUADRANTS = ("tl", "tr", "bl", "br")
ROIS_PATH = Path("rois.json")
RESULTS_PATH = Path("results.json")

# sRGB -> linearna intenziteta (pregledna tabela za 8-bitne vrednosti)
_s = np.arange(256, dtype=np.float64) / 255.0
LIN_LUT = np.where(_s <= 0.04045, _s / 12.92, ((_s + 0.055) / 1.055) ** 2.4).astype(np.float32)
GRAY_WEIGHTS = (0.2126, 0.7152, 0.0722)


# --- nalaganje in parametri ------------------------------------------------


def load_rgb(path):
    """Prebere sliko kot uint8 polje oblike (H, W, 3), z upoštevano EXIF orientacijo."""
    with Image.open(path) as im:
        return np.asarray(ImageOps.exif_transpose(im).convert("RGB"))


def linear_gray(rgb):
    """Vrne (sivina, prežgano): linearizirana svetlost in maska nasičenih pikslov."""
    gray = np.zeros(rgb.shape[:2], dtype=np.float32)
    for channel, weight in enumerate(GRAY_WEIGHTS):
        gray += weight * LIN_LUT[rgb[:, :, channel]]
    clipped = (rgb == 255).any(axis=2)
    return gray, clipped


def params_from_name(name):
    """Izlušči parametre iz imena datoteke (lp = moč laserja v %)."""
    def find(pattern, convert):
        m = re.search(pattern, name)
        return convert(m.group(1)) if m else None

    return {
        "lp": find(r"lp(\d+)", int),
        "v": find(r"\bv(\d+)", int),
        "z": find(r"z(-?[\d.]+)", float),
        "side": find(r"\b(bot|top)\b", str),
    }


# --- geometrija ------------------------------------------------------------


def normalize(rect):
    x0, y0, x1, y1 = rect
    return (min(x0, x1), min(y0, y1), max(x0, x1), max(y0, y1))


def subboxes(rect, pad):
    """Razdeli pravokotnik na 4 kvadrante in vsakega skrči za delež pad na vsaki strani."""
    x0, y0, x1, y1 = normalize(rect)
    xm, ym = (x0 + x1) / 2, (y0 + y1) / 2
    quadrants = {
        "tl": (x0, y0, xm, ym),
        "tr": (xm, y0, x1, ym),
        "bl": (x0, ym, xm, y1),
        "br": (xm, ym, x1, y1),
    }
    shrunk = {}
    for name in QUADRANTS:
        a, b, c, d = quadrants[name]
        w, h = c - a, d - b
        shrunk[name] = (a + pad * w, b + pad * h, c - pad * w, d - pad * h)
    return shrunk


# --- meritev ---------------------------------------------------------------


def measure_box(gray, clipped, box):
    x0, y0, x1, y1 = (int(round(v)) for v in box)
    if x1 - x0 < 2 or y1 - y0 < 2:
        raise ValueError(f"premajhna merilna škatla: {box}")
    patch = gray[y0:y1, x0:x1]
    return {
        "median": float(np.median(patch)),
        "std": float(patch.std()),
        "clip": float(clipped[y0:y1, x0:x1].mean()),
    }


def measure_grid(gray, clipped, rect, pad):
    """Izmeri eno 2x2 mrežo in vrne slovar z vrednostmi kvadrantov in kontrastom."""
    boxes = subboxes(rect, pad)
    values = {name: measure_box(gray, clipped, boxes[name]) for name in QUADRANTS}

    ordered = sorted(QUADRANTS, key=lambda name: values[name]["median"])
    dark, bright = ordered[:2], ordered[2:]
    i_dark = float(np.mean([values[name]["median"] for name in dark]))
    i_bright = float(np.mean([values[name]["median"] for name in bright]))

    # ločljivost parov: razmik med skupinama proti razpršenosti znotraj skupin
    gap = values[bright[0]]["median"] - values[dark[1]]["median"]
    spread = max(
        values[bright[1]]["median"] - values[bright[0]]["median"],
        values[dark[1]]["median"] - values[dark[0]]["median"],
    )

    return {
        "values": values,
        "bright": bright,
        "dark": dark,
        "i_bright": i_bright,
        "i_dark": i_dark,
        "michelson": (i_bright - i_dark) / (i_bright + i_dark) if i_bright + i_dark > 0 else float("nan"),
        "contrast_ratio": i_bright / i_dark if i_dark > 0 else float("inf"),
        "clip_bright": float(np.mean([values[name]["clip"] for name in bright])),
        "clip_dark": float(np.mean([values[name]["clip"] for name in dark])),
        "std_bright": float(np.mean([values[name]["std"] for name in bright])),
        "std_dark": float(np.mean([values[name]["std"] for name in dark])),
        "ambiguous": gap < spread,
    }


# --- rois.json -------------------------------------------------------------


def read_rois():
    if not ROIS_PATH.exists():
        return {}
    data = json.loads(ROIS_PATH.read_text())
    return {
        name: [{"rect": normalize(t["rect"]), "missing": bool(t.get("missing", False))} for t in tiles]
        for name, tiles in data.items()
    }


def write_rois(rois):
    """Zapiše rois.json, vsako mrežo v svojo vrstico zaradi ročnega urejanja."""
    blocks = []
    for name in sorted(rois):
        tiles = ",\n".join(
            "    {{\"rect\": [{}, {}, {}, {}], \"missing\": {}}}".format(
                *(int(round(v)) for v in tile["rect"]),
                "true" if tile["missing"] else "false",
            )
            for tile in rois[name]
        )
        blocks.append(f"  {json.dumps(name)}: [\n{tiles}\n  ]")
    ROIS_PATH.write_text("{\n" + ",\n".join(blocks) + "\n}\n")


# --- interaktivni izbor ----------------------------------------------------


def draw_grid(ax, rect, pad, missing=False, index=None, lines=()):
    """Nariše obrobo mreže, merilne škatle in napise (indeks nad, vrednosti pod mrežo)."""
    color = "#808080" if missing else "#ff3030"
    x0, y0, x1, y1 = normalize(rect)
    ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, ec=color, lw=1.0, ls="--"))
    if not missing:
        for box in subboxes(rect, pad).values():
            ax.add_patch(Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1],
                                   fill=False, ec=color, lw=1.0))
    if index is not None:
        ax.text(x0, y0 - 8, str(index), color=color, fontsize=8, va="bottom", ha="left")
    for i, line in enumerate(("manjka",) if missing else lines):
        ax.text(x0, y1 + 28 + i * 42, line, color=color, fontsize=8, va="top", ha="left")


def select_grids(rgb, name, pad, step):
    """Odpre okno za izbor N_GRIDS mrež. Vrne seznam ploščic ali None ob preklicu."""
    height, width = rgb.shape[:2]
    tiles = []
    state = {"cancelled": False}

    fig, ax = plt.subplots(figsize=(13, 8.5))
    ax.imshow(rgb[::step, ::step], extent=(0, width, height, 0))
    ax.set_axis_off()
    fig.subplots_adjust(0, 0, 1, 0.94)

    def refresh():
        for patch in list(ax.patches):
            patch.remove()
        for text in list(ax.texts):
            text.remove()
        for i, tile in enumerate(tiles):
            draw_grid(ax, tile["rect"], pad, missing=tile["missing"], index=i)
        color = "black" if len(tiles) == N_GRIDS else "firebrick"
        ax.set_title(
            f"{name}   izbranih {len(tiles)}/{N_GRIDS}   "
            "[vleci = nova mreža, m = manjka, u = razveljavi, enter = potrdi, esc = opusti]",
            fontsize=9, color=color,
        )
        fig.canvas.draw_idle()

    def on_select(eclick, erelease):
        if len(tiles) >= N_GRIDS:
            print(f"  že izbranih {N_GRIDS} mrež, najprej razveljavi z 'u'", file=sys.stderr)
            return
        rect = normalize((eclick.xdata, eclick.ydata, erelease.xdata, erelease.ydata))
        tiles.append({"rect": rect, "missing": False})
        refresh()

    def on_key(event):
        if event.key == "u" and tiles:
            tiles.pop()
            refresh()
        elif event.key == "m" and tiles:
            tiles[-1]["missing"] = not tiles[-1]["missing"]
            refresh()
        elif event.key == "enter":
            if len(tiles) == N_GRIDS:
                plt.close(fig)
            else:
                print(f"  potrebnih je točno {N_GRIDS} mrež (izbranih {len(tiles)})", file=sys.stderr)
        elif event.key == "escape":
            state["cancelled"] = True
            plt.close(fig)

    selector = RectangleSelector(  # noqa: F841 (referenca mora živeti do zaprtja okna)
        ax, on_select, useblit=True, button=[1], minspanx=20, minspany=20,
        spancoords="data", interactive=False,
        props=dict(facecolor="none", edgecolor="#30ff30", lw=1.2),
    )
    fig.canvas.mpl_connect("key_press_event", on_key)
    refresh()
    plt.show()

    if state["cancelled"] or len(tiles) != N_GRIDS:
        return None
    return tiles


# --- izhodi ----------------------------------------------------------------


def save_overlay(rgb, name, tiles, pad, results, step):
    """Shrani pomanjšano sliko z označenimi merilnimi območji."""
    height, width = rgb.shape[:2]
    fig, ax = plt.subplots(figsize=(13, 8.5))
    ax.imshow(rgb[::step, ::step], extent=(0, width, height, 0))
    ax.set_axis_off()
    for i, (tile, res) in enumerate(zip(tiles, results)):
        lines = () if res is None else (f"C={res['michelson']:.3f}",
                                        f"R={res['contrast_ratio']:.1f}")
        draw_grid(ax, tile["rect"], pad, missing=tile["missing"], index=i, lines=lines)
    ax.set_title(name, fontsize=9)
    fig.subplots_adjust(0, 0, 1, 0.96)
    path = Path(f"{Path(name).stem}_overlay.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    return path


def write_results(entries):
    """Zapiše results.json kot results[lp][pitch] = {michelson, contrast_ratio, ...}."""
    results = {}
    for entry in entries:
        lp = str(entry["lp"])
        pitch = str(entry["pitch_nm"])
        if pitch in results.get(lp, {}):
            print(f"{entry['file']}: lp={lp}, pitch={pitch} nm že obstaja "
                  f"(iz {results[lp][pitch]['file']}), preskakujem", file=sys.stderr)
            continue
        results.setdefault(lp, {})[pitch] = {k: v for k, v in entry.items()
                                             if k not in ("lp", "pitch_nm")}
    ordered = {lp: {p: results[lp][p] for p in sorted(results[lp], key=int)}
               for lp in sorted(results, key=int)}
    RESULTS_PATH.write_text(json.dumps(ordered, indent=2, ensure_ascii=False) + "\n")


def plot_series(entries, x_key, series_key, y_key, x_label, series_unit, y_label, path):
    series = defaultdict(lambda: defaultdict(list))
    for entry in entries:
        if entry[series_key] is None or entry[x_key] is None or not np.isfinite(entry[y_key]):
            continue
        series[entry[series_key]][entry[x_key]].append(entry[y_key])
    if not series:
        print(f"  ni podatkov za {path}", file=sys.stderr)
        return

    fig, ax = plt.subplots(figsize=(6.5, 4.5))
    keys = sorted(series)
    colors = plt.cm.viridis(np.linspace(0, 0.88, max(len(keys), 1)))
    for color, key in zip(colors, keys):
        xs = sorted(series[key])
        ys = [float(np.mean(series[key][x])) for x in xs]
        ax.plot(xs, ys, "o-", color=color, ms=4, lw=1.2, label=f"{key} {series_unit}")
    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8, ncol=2)
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"  {path}")


def make_plots(entries):
    michelson_label = "Michelsonov kontrast"
    ratio_label = r"razmerje $I_\mathrm{svetlo}/I_\mathrm{temno}$"
    plot_series(entries, "lp", "pitch_nm", "michelson", "moč laserja (%)", "nm",
                michelson_label, Path("contrast_michelson_vs_lp.png"))
    plot_series(entries, "lp", "pitch_nm", "contrast_ratio", "moč laserja (%)", "nm",
                ratio_label, Path("contrast_ratio_vs_lp.png"))
    plot_series(entries, "pitch_nm", "lp", "michelson", "pitch (nm)", "%",
                michelson_label, Path("contrast_michelson_vs_pitch.png"))
    plot_series(entries, "pitch_nm", "lp", "contrast_ratio", "pitch (nm)", "%",
                ratio_label, Path("contrast_ratio_vs_pitch.png"))


# --- glavni tok ------------------------------------------------------------


def process_image(path, rois, args):
    name = path.name
    rgb = None

    if args.reselect or name not in rois:
        if args.compute_only:
            print(f"{name}: ni v {ROIS_PATH}, preskakujem (--compute-only)", file=sys.stderr)
            return []
        print(f"{name}: izberi {N_GRIDS} mrež ...")
        rgb = load_rgb(path)
        selected = select_grids(rgb, name, args.pad, args.step)
        if selected is None:
            print(f"{name}: izbor opuščen, slika preskočena", file=sys.stderr)
            return []
        rois[name] = selected
        write_rois(rois)

    tiles = rois[name]
    if len(tiles) != N_GRIDS:
        print(f"{name}: {len(tiles)} mrež namesto {N_GRIDS}, preskakujem", file=sys.stderr)
        return []

    if rgb is None:
        rgb = load_rgb(path)
    gray, clipped = linear_gray(rgb)
    params = params_from_name(name)

    entries, results = [], []
    for i, tile in enumerate(tiles):
        if tile["missing"]:
            results.append(None)
            continue
        res = measure_grid(gray, clipped, tile["rect"], args.pad)
        results.append(res)
        if res["ambiguous"]:
            print(f"{name} [{i}]: svetla in temna dvojica nista jasno ločeni", file=sys.stderr)
        if res["clip_bright"] > 0.02:
            print(f"{name} [{i}]: {res['clip_bright']:.1%} prežganih pikslov v svetlih kvadratih, "
                  "kontrast je spodnja meja", file=sys.stderr)
        entries.append({
            "lp": params["lp"],
            "pitch_nm": (i + 1) * PITCH_STEP_NM,
            "michelson": res["michelson"],
            "contrast_ratio": res["contrast_ratio"],
            "I_bright": res["i_bright"],
            "I_dark": res["i_dark"],
            **{f"I_{q}": res["values"][q]["median"] for q in QUADRANTS},
            "bright_quadrants": "+".join(sorted(res["bright"])),
            "std_bright": res["std_bright"],
            "std_dark": res["std_dark"],
            "clip_bright": res["clip_bright"],
            "clip_dark": res["clip_dark"],
            "ambiguous": res["ambiguous"],
            "index": i,
            "rect": [int(round(v)) for v in tile["rect"]],
            "file": name,
            "pad": args.pad,
        })

    overlay = save_overlay(rgb, name, tiles, args.pad, results, args.step)
    print(f"  {overlay}")
    return entries


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("images", nargs="*", type=Path, help="slike za obdelavo (privzeto vse *.jpg)")
    parser.add_argument("--pad", type=float, default=0.25,
                        help="delež kvadranta, odrezan na vsaki strani "
                             "(privzeto 0.25, meri se osrednjih 50%%)")
    parser.add_argument("--reselect", action="store_true",
                        help="ponovno izberi mreže, tudi če so že shranjene")
    parser.add_argument("--compute-only", action="store_true",
                        help="brez izbiranja, samo izračun iz rois.json")
    parser.add_argument("--step", type=int, default=4, help="faktor pomanjšanja za prikaz (privzeto 4)")
    args = parser.parse_args()

    if not 0 <= args.pad < 0.5:
        parser.error("--pad mora biti na intervalu [0, 0.5)")

    images = args.images or sorted(Path(".").glob("*.jpg"))
    if not images:
        parser.error("ni slik za obdelavo")

    rois = read_rois()
    entries = []
    for path in images:
        if not path.exists():
            print(f"{path}: datoteke ni, preskakujem", file=sys.stderr)
            continue
        entries.extend(process_image(path, rois, args))

    if not entries:
        print("ni rezultatov", file=sys.stderr)
        return 1

    write_results(entries)
    print(f"  {RESULTS_PATH} ({len(entries)} mrež)")
    make_plots(entries)
    return 0


if __name__ == "__main__":
    sys.exit(main())