"""What is the field correlation made of, and how much of it can matching use?

A puzzle motivates this. At 256 scenes the information-theoretic threshold is
0.29 and the fields correlate at 0.656, so the pairing is identifiable by a
wide margin, yet even at 64 scenes -- where the threshold is 0.51 and the
search space is trivially small -- recovery reaches 14%. Either the theory is
the wrong reference or the statistic is not measuring what the theory means.

The suspect is nuisance structure shared by both fields. Scene similarity has
a component that says only "this scene is busy, so it resembles everything",
and both modalities see it. That component inflates the correlation over all
pairs while carrying no information about which scene is which, because it is
the same for every candidate pairing of a scene to a partner of similar
busyness. Correlated-matching theory assumes the correlation is in the
exchangeable noise, so a nuisance-inflated statistic overstates what the
theory would predict.

This decomposes each field into a degree component -- the rank-one part
predicted by row and column means -- and the residual, and reports the
correlation of each. The residual correlation is the honest input to the
threshold, and if it falls far below 0.656 the recovery gap is explained and
the target moves.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", default="artifacts/vg_5k/natural_pipeline.pt")
    parser.add_argument("--label", default="natural")
    parser.add_argument("--output", default="artifacts/vg_5k/field_anatomy.json")
    return parser.parse_args()


def offdiagonal(matrix: np.ndarray) -> np.ndarray:
    return matrix[~np.eye(len(matrix), dtype=bool)]


def correlate(first: np.ndarray, second: np.ndarray) -> float:
    return float(np.corrcoef(offdiagonal(first), offdiagonal(second))[0, 1])


def degree_part(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Split into the additive row/column-mean model and its residual.

    The two-way additive fit m + r_i + c_j is what a pure busyness effect
    produces: every entry explained by how prominent each of its two scenes
    is, with nothing said about the pair.
    """
    size = len(matrix)
    mask = ~np.eye(size, dtype=bool)
    work = matrix.copy().astype(np.float64)
    np.fill_diagonal(work, np.nan)
    grand = np.nanmean(work)
    rows = np.nanmean(work, axis=1, keepdims=True) - grand
    cols = np.nanmean(work, axis=0, keepdims=True) - grand
    fitted = grand + rows + cols
    residual = np.where(mask, work - fitted, 0.0)
    return np.where(mask, fitted, 0.0), residual


def spectral_strip(matrix: np.ndarray, components: int) -> np.ndarray:
    """Remove the leading eigen-components, a stronger nuisance model."""
    symmetric = (matrix + matrix.T) / 2.0
    values, vectors = np.linalg.eigh(symmetric)
    order = np.argsort(np.abs(values))[::-1]
    keep = order[components:]
    return (vectors[:, keep] * values[keep]) @ vectors[:, keep].T


def main() -> None:
    args = parse_args()
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual = state["visual_field"].double().numpy()
    text = state["text_field"].double().numpy()
    size = len(visual)

    total = correlate(visual, text)
    visual_fit, visual_residual = degree_part(visual)
    text_fit, text_residual = degree_part(text)

    rows = [
        {"component": "raw field", "correlation": total},
        {"component": "degree model only", "correlation": correlate(visual_fit, text_fit)},
        {
            "component": "residual after degree",
            "correlation": correlate(visual_residual, text_residual),
        },
    ]
    for components in (1, 2, 5, 10):
        rows.append(
            {
                "component": f"residual after top-{components} eigen",
                "correlation": correlate(
                    spectral_strip(visual, components), spectral_strip(text, components)
                ),
            }
        )

    # how much of each field's own variance the nuisance model absorbs
    share = {
        "visual_degree_variance_share": float(
            offdiagonal(visual_fit).var() / offdiagonal(visual).var()
        ),
        "text_degree_variance_share": float(
            offdiagonal(text_fit).var() / offdiagonal(text).var()
        ),
    }

    residual_rho = rows[2]["correlation"]
    threshold = float(np.sqrt(4 * np.log(size) / size))
    summary = {
        "protocol": (
            "Both fields split into an additive row/column-mean model and its "
            "residual, then correlated component by component. Hidden pairs "
            "are read only to align the two fields."
        ),
        "label": args.label,
        "size": size,
        "rows": rows,
        **share,
        "it_threshold_at_size": threshold,
        "residual_above_it_threshold": bool(residual_rho > threshold),
        "reading": (
            "If the residual correlation is far below the raw correlation, the "
            "headline statistic is inflated by nuisance structure that carries "
            "no matching information, and the residual is the number that "
            "should be compared against the recovery threshold."
        ),
    }
    for row in rows:
        print(f"{row['component']:<34} rho={row['correlation']:.4f}")
    print(json.dumps(share))
    print(f"IT threshold at N={size}: {threshold:.4f}")
    write_json(args.output, summary)


if __name__ == "__main__":
    main()