"""Data for the bold-line intersection detector (CenterNet style) on v5 canvases.

Ground truth: bold 5 mm lines lie at U = ou + 5i, V = ov + 5j; with a = (U - ou) / 5, b = (V - ov) / 5
from the dense uv map, an intersection is found in every 2 x 2 pixel cell by one Newton step. Nodes
inside the grid area are kept; a border lying on a bold line keeps its nodes, at other borders nodes
closer than MARGIN_TOL_MM are dropped.
Windows: mostly >= 44 of 49 patches on the sheet; EDGE_WINDOW_PROB of them show the sheet border.
Occlusions: 0-2 blobs (blot, blur, noise, glare); the targets stay unchanged.
Targets: 28 x 28 cells of 8 px: heat (Gaussian around each node), offset inside the cell, node flag.
"""

import math

import cv2
import numpy as np

import canvas
import crops

CROP_SIZE = crops.CROP_SIZE
CELL = 8
CELLS = CROP_SIZE // CELL
HEAT_SIGMA = 0.7
EDGE_WINDOW_PROB = 0.3
EDGE_MIN_VALID = 20
MARGIN_TOL_MM = 0.3
OCCLUSION_COUNT_P = (0.2, 0.5, 0.3)
OCCLUSION_AREA_PATCHES = (0.6, 2.2)
OCCLUSION_KINDS = ("blot", "blur", "noise", "glare")
PLAIN_KEEP_EVERY = 3                     # train: every 3rd "plain_1mm" sheet (answer: no nodes)
KIND_CODES = {"red": 0, "dots": 1, "plain_1mm": 2}
CH_HEAT, CH_OFF_X, CH_OFF_Y, CH_POS = range(4)


def canvas_points(uv_mm, valid, texture_info):
    """Bold-line intersections in a canvas. -> [K, 2] x, y (canvas px)."""
    grid = texture_info["grid"]
    if grid["kind"] == "plain_1mm":
        return np.zeros((0, 2))
    ou, ov = grid["bold_line_mm"]
    a = (uv_mm[..., 0].astype(np.float64) - ou) / 5.0
    b = (uv_mm[..., 1].astype(np.float64) - ov) / 5.0
    ok = valid[:-1, :-1] & valid[:-1, 1:] & valid[1:, :-1] & valid[1:, 1:]
    ac = 0.25 * (a[:-1, :-1] + a[:-1, 1:] + a[1:, :-1] + a[1:, 1:])
    bc = 0.25 * (b[:-1, :-1] + b[:-1, 1:] + b[1:, :-1] + b[1:, 1:])
    ax = 0.5 * (a[:-1, 1:] - a[:-1, :-1] + a[1:, 1:] - a[1:, :-1])
    ay = 0.5 * (a[1:, :-1] - a[:-1, :-1] + a[1:, 1:] - a[:-1, 1:])
    bx = 0.5 * (b[:-1, 1:] - b[:-1, :-1] + b[1:, 1:] - b[1:, :-1])
    by = 0.5 * (b[1:, :-1] - b[:-1, :-1] + b[1:, 1:] - b[:-1, 1:])
    det = ax * by - ay * bx
    i, j = np.round(ac), np.round(bc)
    da, db = i - ac, j - bc
    with np.errstate(divide="ignore", invalid="ignore"):
        dx = (by * da - ay * db) / det
        dy = (-bx * da + ax * db) / det
    hit = ok & (np.abs(det) > 1e-9) & (np.abs(dx) <= 0.5) & (np.abs(dy) <= 0.5)
    ys, xs = np.nonzero(hit)
    if len(ys) == 0:
        return np.zeros((0, 2))
    px = xs + 0.5 + dx[ys, xs]
    py = ys + 0.5 + dy[ys, xs]
    ii, jj = i[ys, xs], j[ys, xs]
    u, v = ou + 5.0 * ii, ov + 5.0 * jj
    x0, y0, x1, y1 = grid["area_mm"]
    bold = grid["edges_bold"]
    width_mm, height_mm = texture_info["size_mm"]
    t = MARGIN_TOL_MM
    lo_u = max(x0 - t if bold[0] else x0 + t, t)
    lo_v = max(y0 - t if bold[1] else y0 + t, t)
    hi_u = min(x1 + t if bold[2] else x1 - t, width_mm - t)
    hi_v = min(y1 + t if bold[3] else y1 - t, height_mm - t)
    inside = (u > lo_u) & (u < hi_u) & (v > lo_v) & (v < hi_v)
    keys = np.stack([ii, jj], axis=1)[inside]
    if len(keys) == 0:
        return np.zeros((0, 2))
    # a node on the border of two cells is found twice: average
    _, group = np.unique(keys, axis=0, return_inverse=True)
    group = group.reshape(-1)
    sums = np.zeros((group.max() + 1, 2))
    np.add.at(sums, group, np.stack([px, py], axis=1)[inside])
    return sums / np.bincount(group)[:, None]


def points_to_crop(points, window):
    x0, y0, s = window
    x = (points[:, 0] + 0.5 - x0) / s - 0.5
    y = (points[:, 1] + 0.5 - y0) / s - 0.5
    keep = (x >= -0.5) & (x < CROP_SIZE - 0.5) & (y >= -0.5) & (y < CROP_SIZE - 0.5)
    return np.stack([x, y], axis=1)[keep]


def _blob_mask(rng, size):
    """Soft rotated ellipse with a wobbly border, 0..1."""
    area = rng.uniform(*OCCLUSION_AREA_PATCHES) * crops.PATCH_SIZE ** 2
    aspect = rng.uniform(1.0, 2.5)
    ra = math.sqrt(area * aspect / math.pi)
    rb = ra / aspect
    cx, cy = rng.uniform(0.1, 0.9) * size, rng.uniform(0.1, 0.9) * size
    angle = rng.uniform(0, math.pi)
    soft = rng.uniform(0.8, 3.0)
    reach = int(math.ceil(1.4 * ra + 3 * soft + 2))
    xa, xb = max(0, int(cx) - reach), min(size, int(cx) + reach + 1)
    ya, yb = max(0, int(cy) - reach), min(size, int(cy) + reach + 1)
    yy, xx = np.mgrid[ya:yb, xa:xb].astype(np.float32)
    xr = (xx - cx) * math.cos(angle) + (yy - cy) * math.sin(angle)
    yr = -(xx - cx) * math.sin(angle) + (yy - cy) * math.cos(angle)
    radius = np.sqrt((xr / ra) ** 2 + (yr / rb) ** 2)
    coarse = rng.normal(0, 1, (6, 6)).astype(np.float32)
    wobble = cv2.resize(coarse, (xb - xa, yb - ya), interpolation=cv2.INTER_CUBIC) * rng.uniform(0.1, 0.3)
    mask = np.zeros((size, size), np.float32)
    mask[ya:yb, xa:xb] = (radius + wobble < 1.0).astype(np.float32)
    return cv2.GaussianBlur(mask, (0, 0), soft)


def _smooth_field(rng, size, cells, amplitude):
    coarse = rng.normal(0, 1, (cells, cells)).astype(np.float32)
    return cv2.resize(coarse, (size, size), interpolation=cv2.INTER_CUBIC) * amplitude


def occlude(rng, crop_bgr):
    """0-2 occlusion blobs."""
    size = crop_bgr.shape[0]
    count = int(rng.choice(len(OCCLUSION_COUNT_P), p=OCCLUSION_COUNT_P))
    image = crop_bgr.astype(np.float32)
    for _ in range(count):
        alpha = _blob_mask(rng, size)
        kind = OCCLUSION_KINDS[int(rng.integers(len(OCCLUSION_KINDS)))]
        if kind == "blot":
            color = np.array([[150, 60, 25], [30, 30, 35], [45, 85, 140], [60, 60, 170],
                              [20, 50, 90]][int(rng.integers(5))], np.float32)
            color = color * rng.uniform(0.7, 1.3)
            layer = color[None, None, :] * (1.0 + _smooth_field(rng, size, 12, 0.15))[..., None]
            layer += rng.normal(0, 6, (size, size, 1)).astype(np.float32)
            strength = rng.uniform(0.8, 1.0)
        elif kind == "blur":
            layer = cv2.GaussianBlur(image, (0, 0), rng.uniform(5.0, 12.0))
            strength = 1.0
        elif kind == "noise":
            sigma = rng.uniform(45.0, 100.0)
            grain = rng.normal(0, sigma, (size, size, 3 if rng.random() < 0.5 else 1)).astype(np.float32)
            layer = cv2.GaussianBlur(image, (0, 0), 1.5) + grain
            strength = 1.0
        else:                                   # glare: a saturated core fading out
            level = rng.uniform(235, 255)
            tint = np.array([1.0, rng.uniform(0.97, 1.0), rng.uniform(0.95, 1.0)], np.float32)
            layer = np.full_like(image, level) * tint
            alpha = np.clip(cv2.GaussianBlur(alpha, (0, 0), rng.uniform(3.0, 8.0)) * 1.4, 0, 1)
            strength = rng.uniform(0.9, 1.0)
        a = (alpha * strength)[..., None]
        image = image * (1 - a) + layer * a
    return np.clip(image, 0, 255).astype(np.uint8)


def make_targets(points_crop):
    """-> [28, 28, 4] float32: heat, offset x, offset y, node flag."""
    target = np.zeros((CELLS, CELLS, 4), np.float32)
    if len(points_crop) == 0:
        return target
    pc = (points_crop + 0.5) / CELL
    cell = np.clip(np.floor(pc).astype(int), 0, CELLS - 1)
    offset = np.clip(pc - cell, 0.0, 1.0 - 1e-6)
    grid = np.arange(CELLS)
    heat = np.zeros((CELLS, CELLS), np.float32)
    for (cx, cy) in cell:
        heat = np.maximum(heat, np.exp(-((grid[None, :] - cx) ** 2 + (grid[:, None] - cy) ** 2) / (2 * HEAT_SIGMA ** 2)))
    for (cx, cy), (ox, oy) in zip(cell, offset):
        heat[cy, cx] = 1.0
        target[cy, cx, CH_OFF_X] = ox
        target[cy, cx, CH_OFF_Y] = oy
        target[cy, cx, CH_POS] = 1.0
    target[..., CH_HEAT] = heat
    return target


def train_selection(directory, indices):
    """No "dots" sheets (their dots are invisible at 3.5-7.5 px/mm), every PLAIN_KEEP_EVERY-th "plain_1mm" one."""
    kind = {i: canvas.read_texture_info(directory, i)["grid"]["kind"] for i in indices}
    keep_plain = set([i for i in indices if kind[i] == "plain_1mm"][::PLAIN_KEEP_EVERY])
    return [i for i in indices if kind[i] == "red" or i in keep_plain]


def make_sample(rng, directory, index, cache):
    """-> (crop BGR uint8, target [28, 28, 4], grid kind code)."""
    image, uv_mm, valid = canvas.read_sample(directory, index)
    texture = canvas.read_texture_info(directory, index)
    if rng.random() < EDGE_WINDOW_PROB:
        window = crops.choose_window(rng, uv_mm, valid, image.shape[0], EDGE_MIN_VALID, crops.MIN_VALID_PATCHES)
    else:
        window = crops.choose_window(rng, uv_mm, valid, image.shape[0])
    crop = occlude(rng, crops.crop_image(image, window))
    if index not in cache:
        cache[index] = canvas_points(uv_mm, valid, texture)
    return crop, make_targets(points_to_crop(cache[index], window)), KIND_CODES[texture["grid"]["kind"]]


class PointDataset:
    """train: new random window and occlusions every call; otherwise one seeded window per index."""

    def __init__(self, directory, indices, train=True, fixed_seed=777):
        self.directory, self.indices, self.train, self.fixed_seed = directory, list(indices), train, fixed_seed
        self.rng = None
        self.cache = {}

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, position):
        import torch
        index = self.indices[position]
        if self.train:
            if self.rng is None:
                self.rng = np.random.default_rng(torch.initial_seed() % (2 ** 63))
            rng = self.rng
        else:
            rng = np.random.default_rng([self.fixed_seed, index])
        crop, target, kind = make_sample(rng, self.directory, index, self.cache)
        rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
        return torch.from_numpy(np.ascontiguousarray(rgb.transpose(2, 0, 1))), torch.from_numpy(target), torch.tensor(kind)


def decode(heat_prob, offsets, threshold):
    """3 x 3 local maxima above threshold -> [K, 2] crop px."""
    padded = np.pad(heat_prob, 1, constant_values=-1.0)
    neighbours = np.max(np.stack([padded[dy:dy + CELLS, dx:dx + CELLS] for dy in range(3) for dx in range(3)]), axis=0)
    ys, xs = np.nonzero((heat_prob >= neighbours) & (heat_prob >= threshold))
    return np.stack([(xs + offsets[ys, xs, 0]) * CELL - 0.5, (ys + offsets[ys, xs, 1]) * CELL - 0.5], axis=1)


def target_points(target):
    ys, xs = np.nonzero(target[..., CH_POS] > 0.5)
    return np.stack([(xs + target[ys, xs, CH_OFF_X]) * CELL - 0.5, (ys + target[ys, xs, CH_OFF_Y]) * CELL - 0.5], axis=1)


def match_count(predicted, truth, radius):
    """Greedy one-to-one matching by distance -> number of matched pairs."""
    if len(predicted) == 0 or len(truth) == 0:
        return 0
    distance = np.hypot(predicted[:, None, 0] - truth[None, :, 0], predicted[:, None, 1] - truth[None, :, 1])
    pi, ti = np.nonzero(distance <= radius)
    used_p, used_t = set(), set()
    for k in np.argsort(distance[pi, ti]):
        if pi[k] not in used_p and ti[k] not in used_t:
            used_p.add(pi[k])
            used_t.add(ti[k])
    return len(used_p)


def f1_scores(heat_prob, offsets, targets, kinds, thresholds, radius):
    """Precision / recall / F1 on "red" crops per threshold and the best F1."""
    result, best = {}, -1.0
    for threshold in thresholds:
        tp = fp = fn = 0
        for n in np.nonzero(kinds == KIND_CODES["red"])[0]:
            predicted = decode(heat_prob[n], offsets[n], threshold)
            truth = target_points(targets[n])
            m = match_count(predicted, truth, radius)
            tp, fp, fn = tp + m, fp + len(predicted) - m, fn + len(truth) - m
        precision, recall = tp / max(tp + fp, 1), tp / max(tp + fn, 1)
        f1 = 2 * precision * recall / max(precision + recall, 1e-12)
        result["t{:02d}".format(int(round(threshold * 100)))] = (precision, recall, f1)
        best = max(best, f1)
    result["best_f1"] = best
    return result