"""Photo scale (px/mm) from detector trials at several assumed scales.

A square crop around a point is resized to 5.5 px/mm under each assumed scale; the detector finds
nodes and they are linked into a lattice. A right guess gives a large lattice with a 27.5 px (5 mm)
step. The "5 mm cell" check (four thin 1 mm lines inside a cell -> 5th harmonic of the brightness
profile) rejects the x5 (1 mm grid taken for 5 mm) and x2 (10 mm) errors.
"""

import cv2
import numpy as np

from . import lattice, nets

TARGET_PX_PER_MM = 5.5
STEP_PX = 5.0 * TARGET_PX_PER_MM
CENTER_CROP = 448
TRIAL_PX_PER_MM = (2.0, 3.3, 5.5, 9.0, 15.0, 25.0, 40.0)
PROFILE_SAMPLES = 100
MIN_CELLS = 20
MIN_FIVE_10MM = 0.5
MIN_CELL_PX = 12.0
MIN_FIVE = 1.2
SCALE_TRIES = 3
MAX_SHEET_MM = 600.0


def resize(image, factor):
    """INTER_AREA for downscaling (no moire on thin grid lines), cubic for upscaling."""
    if abs(factor - 1.0) < 1e-3:
        return image
    return cv2.resize(image, None, fx=factor, fy=factor,
                      interpolation=cv2.INTER_AREA if factor < 1 else cv2.INTER_CUBIC)


def five_score(image, quads):
    """5th harmonic of the mean brightness profile across lattice cells / mean of the 4th and 6th."""
    green = image[..., 1].astype(np.float32)
    t = np.arange(PROFILE_SAMPLES) / PROFILE_SAMPLES
    profiles = []
    for a, b, c, d in ((0, 1, 3, 2), (0, 3, 1, 2)):
        for offset in (0.35, 0.5, 0.65):
            start = quads[:, a] + offset * (quads[:, c] - quads[:, a])
            end = quads[:, b] + offset * (quads[:, d] - quads[:, b])
            line = start[:, None] + t[None, :, None] * (end - start)[:, None]
            profiles.append(cv2.remap(green, line[..., 0].astype(np.float32), line[..., 1].astype(np.float32),
                                      cv2.INTER_LINEAR))
    mean = np.concatenate(profiles).mean(0)
    amplitude = np.abs(np.fft.rfft(mean - mean.mean()))
    return float(amplitude[5] / max(1e-6, 0.5 * (amplitude[4] + amplitude[6])))


def center_trial(models, image, px_per_mm, cx, cy, min_px_per_mm=0.0):
    """One trial around (cx, cy) assuming px_per_mm. score 0 means unusable."""
    h, w = image.shape[:2]
    factor = TARGET_PX_PER_MM / px_per_mm
    half = CENTER_CROP / factor / 2
    x0, x1 = int(max(0, cx - half)), int(min(w, cx + half))
    y0, y1 = int(max(0, cy - half)), int(min(h, cy + half))
    crop = resize(image[y0:y1, x0:x1], factor)
    found = nets.detector_net(models, crop)
    if len(found) < 8:
        return {"score": 0}
    xy = found[:, :2]
    index, component, right, down = lattice.index_mutual(xy, STEP_PX)
    members = np.nonzero(component == np.argmax(np.bincount(component)))[0]
    step_x, step_y, nr, nd = lattice.lattice_stats(xy, right, down, members)
    step = np.nanmean([step_x, step_y])
    closeness = np.exp(-0.5 * (np.log(step / STEP_PX) / 0.2) ** 2) if np.isfinite(step) else 0.0
    quads = lattice.cell_quads(xy, index, members)
    five = five_score(image, (quads + 0.5) / factor - 0.5 + np.array([x0, y0])) if len(quads) >= 4 else 0.0
    if np.sqrt(abs(px_per_mm * step_x / STEP_PX * px_per_mm * step_y / STEP_PX)) < min_px_per_mm:
        closeness = 0.0
    return {"score": float(min(nr, nd) * closeness), "cells": int(len(quads)), "five": five,
            "px_per_mm_x": px_per_mm * step_x / STEP_PX, "px_per_mm_y": px_per_mm * step_y / STEP_PX}


def find_scale_at(models, image, cx, cy, min_px_per_mm=0.0):
    """Best trial around (cx, cy) with a "verified" flag, or None.
    Usable trials are grouped (+-15 %); 10 mm lattices are dropped; a group is verified when its
    cells look like 5 mm cells. The winner is refined by one more trial at the found scale."""
    trial = lambda p: center_trial(models, image, p, cx, cy, min_px_per_mm)
    trials = [trial(p) for p in TRIAL_PX_PER_MM]
    estimate = lambda t: float(np.sqrt(t["px_per_mm_x"] * t["px_per_mm_y"]))
    usable = [t for t in trials if t["score"] > 0 and t.get("cells", 0) >= MIN_CELLS
              and abs(np.log(t["px_per_mm_x"] / t["px_per_mm_y"])) < np.log(1.25)]
    if not usable:
        return None
    groups = []
    for t in sorted(usable, key=lambda t: -t["cells"]):
        for g in groups:
            if abs(np.log(estimate(t) / g["px_per_mm"])) < np.log(1.15):
                g["trials"].append(t)
                g["weight"] += t["cells"]
                g["five"] = max(g["five"], t["five"])
                break
        else:
            groups.append({"px_per_mm": estimate(t), "trials": [t], "weight": t["cells"], "five": t["five"]})
    groups = [g for g in groups if g["five"] >= MIN_FIVE_10MM] or groups
    groups.sort(key=lambda g: -g["weight"])
    top = groups[0]["weight"]
    verified = [g for g in groups if g["five"] >= MIN_FIVE and 5 * g["px_per_mm"] >= MIN_CELL_PX
                and g["weight"] >= 0.25 * top]
    if verified:
        best = max(verified, key=lambda g: g["weight"] * min(g["five"], 10.0))
    else:
        best = groups[0]
        for g in groups[1:]:
            ratio = best["px_per_mm"] / g["px_per_mm"]
            if 1.7 < ratio < 2.4 and g["weight"] > 0.5 * best["weight"]:
                best = g
    seed = max(best["trials"], key=lambda t: t["cells"])
    refined = trial(estimate(seed))
    agrees = refined["score"] > 0 and refined.get("cells", 0) >= MIN_CELLS and \
        abs(np.log(estimate(refined) / estimate(seed))) < np.log(1.15)
    return dict(refined if agrees else seed, verified=bool(verified))


def find_scale_in_mask(models, image, mask):
    """Scale at the deepest mask point, then at up to SCALE_TRIES - 1 other deep points.
    -> (px/mm or None, (cx, cy) or None, tried points, verified)."""
    dist = cv2.distanceTransform(mask.astype(np.uint8), cv2.DIST_L2, 5)
    min_ppm = max(image.shape[:2]) / MAX_SHEET_MM
    tried = []
    for _ in range(SCALE_TRIES):
        cy, cx = np.unravel_index(int(np.argmax(dist)), dist.shape)
        if dist[cy, cx] <= 0:
            break
        centre = find_scale_at(models, image, cx, cy, min_ppm)
        ppm = None if centre is None else float(np.sqrt(centre["px_per_mm_x"] * centre["px_per_mm_y"]))
        tried.append([int(cx), int(cy), None if ppm is None else round(ppm, 3)])
        if centre is not None:
            return ppm, (int(cx), int(cy)), tried, bool(centre["verified"])
        r = int(max(20, 0.15 * max(image.shape[:2])))
        cv2.circle(dist, (int(cx), int(cy)), r, 0, -1)
    return None, None, tried, False