"""Choose stroke angles that make nearby camera-grid neighborhoods distinct."""

import math
import random

import numpy as np

from camera_grid_stroke_geometry import stroke_angle


ANGLE_COUNT = 12
OFFSETS = ((1, 0), (0, 1), (1, 1), (1, -1),
           (2, 0), (0, 2), (2, 2), (2, -2))


def _cell_template(angle, width_ratio, length_ratio, resolution):
    sample = 4
    side = resolution * sample
    axis = (np.arange(side, dtype=np.float64) + 0.5) / side - 0.5
    x, y = np.meshgrid(axis, axis)
    along = x * math.cos(angle) + y * math.sin(angle)
    across = -x * math.sin(angle) + y * math.cos(angle)
    half_body = (length_ratio - width_ratio) / 2
    outside = np.maximum(np.abs(along) - half_body, 0)
    binary = (outside * outside + across * across <= (width_ratio / 2) ** 2)
    return binary.reshape(resolution, sample, resolution, sample).mean(axis=(1, 3))


def _angle_distances(width_ratio, length_ratio):
    angles = np.arange(ANGLE_COUNT) * math.pi / ANGLE_COUNT
    distances = np.zeros((ANGLE_COUNT, ANGLE_COUNT), dtype=np.float64)
    for resolution in (4, 8):
        templates = np.stack([_cell_template(a, width_ratio, length_ratio, resolution)
                              for a in angles])
        diff = templates[:, None] - templates[None, :]
        distances += np.mean(diff * diff, axis=(2, 3)) / 2
    peak = distances.max()
    if peak > 0:
        distances /= peak
    return angles, distances


def _neighborhood_pairs(nx, ny):
    windows_a = []
    windows_b = []
    incident = [[] for _ in range(nx * ny)]
    for y in range(1, ny - 1):
        for x in range(1, nx - 1):
            for dx, dy in OFFSETS:
                xx, yy = x + dx, y + dy
                if not (1 <= xx < nx - 1 and 1 <= yy < ny - 1):
                    continue
                a = [(y + v) * nx + x + u for v in (-1, 0, 1)
                     for u in (-1, 0, 1)]
                b = [(yy + v) * nx + xx + u for v in (-1, 0, 1)
                     for u in (-1, 0, 1)]
                index = len(windows_a)
                windows_a.append(a)
                windows_b.append(b)
                for cell in set(a + b):
                    incident[cell].append(index)
    return (np.asarray(windows_a, dtype=np.int32),
            np.asarray(windows_b, dtype=np.int32),
            [np.asarray(indices, dtype=np.int32) for indices in incident])


def _pair_distances(assignments, distances, windows_a, windows_b):
    return distances[assignments[windows_a], assignments[windows_b]].mean(axis=1)


def optimize_angles(nx, ny, seed, width_ratio=0.42, length_ratio=0.84,
                    proposals=100000):
    """Optimize 3x3 neighborhood separation within a two-cell search radius.

    The 4/8-pixel camera templates are a cheap screening model. The rendered
    optical-flow probe remains the acceptance test for a candidate pattern.
    """
    if nx < 5 or ny < 5 or proposals < 0:
        raise ValueError("require at least 5x5 cells and nonnegative proposals")
    angles, distances = _angle_distances(width_ratio, length_ratio)
    wa, wb, incident = _neighborhood_pairs(nx, ny)
    assignments = np.asarray([
        int(round(stroke_angle(x, y, seed) * ANGLE_COUNT / math.pi)) % ANGLE_COUNT
        for y in range(ny) for x in range(nx)
    ], dtype=np.int16)
    initial = _pair_distances(assignments, distances, wa, wb)
    pair_loss = np.exp(-6 * initial)
    rng = random.Random(seed ^ 0x71C3A5)
    accepted = 0
    for _ in range(proposals):
        cell = rng.randrange(nx * ny)
        affected = incident[cell]
        if not len(affected):
            continue
        previous = assignments[cell]
        candidate = rng.randrange(ANGLE_COUNT - 1)
        if candidate >= previous:
            candidate += 1
        assignments[cell] = candidate
        trial = _pair_distances(assignments, distances, wa[affected], wb[affected])
        trial_loss = np.exp(-6 * trial)
        if trial_loss.sum() < pair_loss[affected].sum() - 1e-12:
            pair_loss[affected] = trial_loss
            accepted += 1
        else:
            assignments[cell] = previous
    final = _pair_distances(assignments, distances, wa, wb)
    report = {
        "pairs": len(initial),
        "proposals": proposals,
        "accepted": accepted,
        "initial_distance_p10": float(np.percentile(initial, 10)),
        "final_distance_p10": float(np.percentile(final, 10)),
        "initial_mean_loss": float(np.exp(-6 * initial).mean()),
        "final_mean_loss": float(pair_loss.mean()),
    }
    return angles[assignments].reshape(ny, nx), report