"""Upper-bound set battery: connected-component sprites as vision sets.

The learned towers have not yet produced object states, which leaves two
hypotheses entangled: the set-kernel machinery could be wrong, or only
the towers could be short. This battery separates them. On this world
the background is flat, so connected bright components ARE the objects;
per-component centered sprites are model-free object states of the
minimal-world-knowledge class (like the color anchors on natural data:
declared, fixed, no learning). If set-kernel relation fields built from
these pass the gate and support recovery, the machinery is validated and
the remaining gap is exactly "an SSL objective that discovers objects".

Text sets are the per-group phrase states of the set battery. Hidden
pairs score orderings only.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from scipy import ndimage
from tqdm import tqdm

from .common import read_json, seed_everything, write_json
from .manifold_gate import standardize_relation
from .ricci_control import run_gates
from .synth_set_battery import (
    parse_group_phrases,
    set_similarity_field,
    text_group_sets,
)
from .synth_towers import load_image


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="artifacts/synth_v0")
    parser.add_argument("--text-tower", default="artifacts/synth_v0/text_tower.pt")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=512)
    parser.add_argument("--sprite-window", type=int, default=56)
    parser.add_argument("--merge-distance", type=float, default=30.0)
    parser.add_argument(
        "--features", choices=["pixels", "descriptors", "onehot"], default="descriptors"
    )
    parser.add_argument("--text-mode", choices=["lm", "bow"], default="bow")
    parser.add_argument(
        "--kernel", choices=["matching", "moment"], default="matching",
        help="moment: symmetric-tensor (Fock) set kernel, no matching step",
    )
    parser.add_argument("--vision-views", type=int, default=4)
    parser.add_argument("--random-perms", type=int, default=300)
    parser.add_argument("--descent-restarts", type=int, default=5)
    parser.add_argument("--descent-max-steps", type=int, default=200000)
    parser.add_argument(
        "--descent-objective", default="mse", choices=["mse", "m30_total"]
    )
    parser.add_argument("--descent-verify-top", type=int, default=64)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--seed", type=int, default=20260731)
    parser.add_argument(
        "--output", default="artifacts/synth_v0/cc_battery_gate.json"
    )
    return parser.parse_args()


def component_sprites(
    image: torch.Tensor, window: int, merge_distance: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """Centered sprites of bright connected components, ring-merged.

    Same-group objects sit on a small ring; merging nearby components
    reassembles groups so a sprite carries multiplicity as pattern.
    """
    array = image.permute(1, 2, 0).numpy()
    background = np.median(array.reshape(-1, 3), axis=0)
    foreground = (np.abs(array - background).sum(-1) > 0.12)
    labels, count = ndimage.label(foreground)
    if count == 0:
        return torch.zeros(1, 3 * window * window), torch.ones(1)
    centers = np.array(ndimage.center_of_mass(foreground, labels, range(1, count + 1)))
    sizes = ndimage.sum(foreground, labels, range(1, count + 1))
    # Merge components whose centers are close (ring members).
    parent = list(range(count))

    def find(a: int) -> int:
        while parent[a] != a:
            parent[a] = parent[parent[a]]
            a = parent[a]
        return a

    for a in range(count):
        for b in range(a + 1, count):
            if np.linalg.norm(centers[a] - centers[b]) < merge_distance:
                parent[find(a)] = find(b)
    groups: dict[int, list[int]] = {}
    for a in range(count):
        groups.setdefault(find(a), []).append(a)

    height, width = foreground.shape
    half = window // 2
    padded = np.pad(array, ((half, half), (half, half), (0, 0)))
    padded_mask = np.pad(foreground, half)
    sprites, weights = [], []
    for members in groups.values():
        member_mask = np.isin(labels, [m + 1 for m in members])
        mass = float(member_mask.sum())
        ys, xs = np.nonzero(member_mask)
        cy, cx = int(ys.mean()), int(xs.mean())
        patch = padded[cy : cy + window, cx : cx + window].copy()
        mask_patch = padded_mask[cy : cy + window, cx : cx + window]
        patch[~mask_patch] = 0.0
        sprites.append(torch.from_numpy(patch).float().flatten())
        weights.append(mass)
    weights = torch.tensor(weights)
    return torch.stack(sprites), weights / weights.sum().clamp_min(1e-8)


HUE_CENTERS = {
    "red": 0.0, "orange": 30.0, "yellow": 60.0, "green": 120.0,
    "cyan": 180.0, "blue": 220.0, "purple": 275.0, "pink": 330.0,
}


def component_descriptors(
    image: torch.Tensor, merge_distance: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """Rotation-invariant group descriptors: color classes, log-area,
    member multiplicity, and normalized single-member area."""
    import colorsys

    array = image.permute(1, 2, 0).numpy()
    background = np.median(array.reshape(-1, 3), axis=0)
    foreground = np.abs(array - background).sum(-1) > 0.12
    labels, count = ndimage.label(foreground)
    if count == 0:
        return torch.zeros(1, 19), torch.ones(1)
    centers = np.array(
        ndimage.center_of_mass(foreground, labels, range(1, count + 1))
    )
    parent = list(range(count))

    def find(a: int) -> int:
        while parent[a] != a:
            parent[a] = parent[parent[a]]
            a = parent[a]
        return a

    for a in range(count):
        for b in range(a + 1, count):
            if np.linalg.norm(centers[a] - centers[b]) < merge_distance:
                parent[find(a)] = find(b)
    groups: dict[int, list[int]] = {}
    for a in range(count):
        groups.setdefault(find(a), []).append(a)
    descriptors, weights = [], []
    total = foreground.size
    for members in groups.values():
        member_mask = np.isin(labels, [m + 1 for m in members])
        pixels = array[member_mask]
        mean_rgb = pixels.mean(0)
        h, s_, v = colorsys.rgb_to_hsv(*mean_rgb.tolist())
        hue = h * 360.0
        vector = np.zeros(16)
        if v > 0.75 and s_ < 0.25:
            vector[8] = 1.0  # white
        elif v < 0.25:
            vector[9] = 1.0  # black
        elif s_ < 0.25:
            vector[10] = 1.0  # gray
        elif abs(hue - 30.0) < 25.0 and v < 0.55:
            vector[11] = 1.0  # brown
        else:
            for index, center in enumerate(HUE_CENTERS.values()):
                distance = min(abs(hue - center), 360.0 - abs(hue - center))
                if distance < 25.0:
                    vector[index] = 1.0
                    break
        member_count = len(members)
        area = float(member_mask.sum()) / total
        vector[12] = np.log(area + 1e-6) / 6.0
        vector[13] = (member_count - 1) / 3.0
        vector[14] = np.log(area / member_count + 1e-6) / 6.0
        vector[15] = 1.0
        # Rotation-invariant shape features of the largest single member:
        # compactness, convexity, and inertia eccentricity separate the
        # six shapes without orientation.
        largest = max(members, key=lambda m: (labels == m + 1).sum())
        single = labels == largest + 1
        area_px = float(single.sum())
        eroded = ndimage.binary_erosion(single)
        perimeter = float((single & ~eroded).sum())
        compactness = 4.0 * np.pi * area_px / max(perimeter, 1.0) ** 2
        ys, xs = np.nonzero(single)
        ys = ys - ys.mean(); xs = xs - xs.mean()
        cov = np.cov(np.stack([xs, ys])) + 1e-6 * np.eye(2)
        eigenvalues = np.linalg.eigvalsh(cov)
        eccentricity = float(1.0 - eigenvalues[0] / eigenvalues[1])
        hull_span = (xs.max() - xs.min() + 1) * (ys.max() - ys.min() + 1)
        boxfill = area_px / max(hull_span, 1.0)
        shape_vector = np.array([compactness, eccentricity, boxfill])
        vector = np.concatenate([vector, shape_vector])
        descriptors.append(torch.tensor(vector, dtype=torch.float32))
        weights.append(float(member_mask.sum()))
    weights = torch.tensor(weights)
    return torch.stack(descriptors), weights / weights.sum().clamp_min(1e-8)


def onehot_descriptors(
    raw_sets: list[torch.Tensor],
) -> list[torch.Tensor]:
    """Factor-mirrored one-hot recoding of descriptor sets.

    Sizes are binned by corpus terciles of single-member log-area with the
    middle bin unmarked, mirroring the text side where medium size has no
    word; shapes are k-means clusters of the rotation-invariant shape
    features. Both statistics come from the evaluated corpus itself,
    unimodally. Output channels mirror the bag-of-words support: color
    (12), small/large (2), count (4), shape cluster (6).
    """
    from sklearn.cluster import KMeans

    all_groups = torch.cat(raw_sets)
    # The renderer shrinks radii with member count, so raw single-member
    # area confounds size class with multiplicity; residualize log-area on
    # member count before binning (unimodal statistics).
    counts_all = (all_groups[:, 13] * 3.0).round().clamp(0, 3)
    area_all = all_groups[:, 14]
    count_means = {}
    for value in (0.0, 1.0, 2.0, 3.0):
        chosen = counts_all == value
        count_means[value] = float(area_all[chosen].mean()) if chosen.any() else 0.0
    adjusted_all = area_all - torch.tensor(
        [count_means[float(v)] for v in counts_all]
    )
    low, high = adjusted_all.quantile(1.0 / 3.0), adjusted_all.quantile(2.0 / 3.0)
    shape_features = all_groups[:, 16:19].numpy()
    clusters = KMeans(n_clusters=6, n_init=10, random_state=0).fit(shape_features)
    recoded = []
    for groups in raw_sets:
        vectors = torch.zeros(len(groups), 24)
        vectors[:, :12] = groups[:, :12]
        counts_here = (groups[:, 13] * 3.0).round().clamp(0, 3)
        adjusted = groups[:, 14] - torch.tensor(
            [count_means[float(v)] for v in counts_here]
        )
        vectors[:, 12] = (adjusted <= low).float()   # small
        vectors[:, 13] = (adjusted >= high).float()  # large
        counts = (groups[:, 13] * 3.0).round().long().clamp(0, 3)
        vectors[torch.arange(len(groups)), 14 + counts] = 1.0
        labels = clusters.predict(groups[:, 16:19].numpy())
        vectors[torch.arange(len(groups)), 18 + labels] = 1.0
        recoded.append(vectors)
    return recoded


def moment_field(sets: list[torch.Tensor]) -> torch.Tensor:
    """Symmetric-tensor (second-quantized) set kernel, matching-free.

    phi(S) concatenates the degree-1 and degree-2 moments of the set; the
    field is the Gram matrix of normalized phi. No assignment step, so no
    matching-value or set-size bias can enter.
    """
    phis = []
    for members in sets:
        m1 = members.mean(0)
        m2 = (members[:, :, None] * members[:, None, :]).mean(0).flatten()
        phi = torch.cat([m1, m2])
        phis.append(phi / phi.norm().clamp_min(1e-9))
    stacked = torch.stack(phis)
    return stacked @ stacked.T


def phrase_bow_sets(
    rows: list[int], captions: list[list[str]], vocabulary: list[str]
) -> list[torch.Tensor]:
    index = {word: i for i, word in enumerate(vocabulary)}
    sets = []
    for row in rows:
        phrases = parse_group_phrases(captions[row][0])
        vectors = torch.zeros(len(phrases), len(index))
        for p, phrase in enumerate(phrases):
            for token in phrase.split():
                if token in index:
                    vectors[p, index[token]] += 1.0
        sets.append(F.normalize(vectors, dim=-1))
    return sets


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    manifest = read_json(Path(args.data_dir, "manifest.json"))
    captions = read_json(Path(args.data_dir, "captions.json"))["captions"]
    rows = manifest[args.split][: args.samples]
    image_dir = Path(manifest["image_dir"])

    per_view_fields = []
    set_sizes = []
    for view in range(args.vision_views):
        vision_sets = []
        raw_sets = []
        for row in tqdm(rows, desc=f"cc sprites v{view}"):
            image = load_image(image_dir / f"scene{row:06d}_v{view}.png")
            if args.features in ("descriptors", "onehot"):
                sprites, _ = component_descriptors(image, args.merge_distance)
            else:
                sprites, _ = component_sprites(
                    image, args.sprite_window, args.merge_distance
                )
            raw_sets.append(sprites)
            if view == 0:
                set_sizes.append(len(sprites))
        if args.features == "onehot":
            vision_sets = [
                F.normalize(v, dim=-1) for v in onehot_descriptors(raw_sets)
            ]
        else:
            vision_sets = [F.normalize(v, dim=-1) for v in raw_sets]
        if args.kernel == "moment":
            per_view_fields.append(moment_field(vision_sets))
        else:
            per_view_fields.append(set_similarity_field(vision_sets))

    if args.text_mode == "bow":
        text_sets = phrase_bow_sets(rows, captions, manifest["vocabulary"])
    else:
        text_sets = text_group_sets(rows, captions, args)
    if args.kernel == "moment":
        text_field = moment_field(text_sets)
    else:
        text_field = set_similarity_field(text_sets)
    # Mass weighting inside the matching corrupts similarity grading
    # (0.16 vs 0.45 against soft truth); match unweighted. Averaging the
    # per-view fields cancels segmentation errors across resampled layouts.
    visual_field = torch.stack(per_view_fields).mean(0)
    visual_channels = standardize_relation(visual_field.double())[0][None]
    text_channels = standardize_relation(text_field.double())[0][None]
    generator = torch.Generator().manual_seed(args.seed)
    report = {
        "protocol": (
            "Vision sets are model-free connected-component sprites "
            "(declared minimal world knowledge); text sets are per-group "
            "phrase states. Hidden pairs score orderings only."
        ),
        "split": args.split,
        "samples": len(rows),
        "mean_vision_set_size": float(np.mean(set_sizes)),
        **run_gates(text_channels, visual_channels, args, generator),
    }
    verdict = {
        "true_z_mse": report["gate_a"]["random"]["mse"]["true_z"],
        "improving_fraction": report["gate_b"]["improving_fraction"],
        "descent_keeps": report["descent_from_true"]["final_accuracy"],
        "true_mse": report["gate_a"]["true"]["mse"],
        "best_random_descent": min(
            (r["final_objective"] for r in report["descent_from_random"]),
            default=None,
        ),
    }
    verdict["counterfeit_found"] = bool(
        verdict["best_random_descent"] is not None
        and verdict["best_random_descent"] < verdict["true_mse"]
    )
    report["verdict"] = verdict
    print(json.dumps({"verdict": verdict}))
    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()