"""Is the effective rank of the shared field the binding variable, not its magnitude?

Stripping ten eigen-directions costs the synthetic fields that recover almost
nothing -- 0.929 to 0.870 -- and costs the natural fields almost everything --
0.656 to 0.080. The natural shared signal is therefore concentrated in about
ten directions while the synthetic one is spread across the spectrum. That is
a correlation between rank and recovery, and this makes it causal by
intervention: take the synthetic fields that do recover, project them to
successively lower rank, and watch recovery as a function of rank while the
correlation between the two truncated fields stays high.

If recovery dies at low rank while the correlation is still far above the
threshold, then rank is the binding variable and the headline correlation is
the wrong target. The theory agrees in advance: correlated-matching thresholds
assume the agreement lives in exchangeable full-rank noise, so N^2/2 entries
each carry an independent constraint. A rank-r shared component supplies about
rN, and at 256 scenes with r near ten that is an eighth of what the threshold
calculation assumes.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json
from .spectral_match import grampa
from .synth_fast_gate import ClosedFormEnergy, all_swaps, steepest_descent
from .synth_triangle_gate import standardized


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", default="artifacts/synth_v1/fields_tier0_ws_256.pt")
    parser.add_argument("--label", default="synthetic-watershed")
    parser.add_argument("--trials", type=int, default=3)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--output", default="artifacts/synth_v1/rank_causal.json")
    return parser.parse_args()


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


def normalise(matrix: np.ndarray) -> np.ndarray:
    values = offdiagonal(matrix)
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


def truncate(matrix: np.ndarray, rank: int) -> np.ndarray:
    """Keep the leading `rank` eigen-components of the symmetrised field."""
    symmetric = (matrix + matrix.T) / 2.0
    values, vectors = np.linalg.eigh(symmetric)
    order = np.argsort(np.abs(values))[::-1][:rank]
    return (vectors[:, order] * values[order]) @ vectors[:, order].T


def effective_rank(matrix: np.ndarray) -> float:
    """Participation ratio of the eigen-spectrum: how many directions carry it."""
    values = np.abs(np.linalg.eigvalsh((matrix + matrix.T) / 2.0))
    weights = values / values.sum()
    weights = weights[weights > 0]
    return float(np.exp(-(weights * np.log(weights)).sum()))


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

    rows = []
    ranks = [4, 8, 16, 32, 64, 128, size]
    for rank in ranks:
        visual = normalise(truncate(visual_full, rank) if rank < size else visual_full)
        text = normalise(truncate(text_full, rank) if rank < size else text_full)
        mask = ~np.eye(size, dtype=bool)
        rho = float(np.corrcoef(visual[mask], text[mask])[0, 1])

        accuracies = []
        for trial in range(args.trials):
            generator = np.random.default_rng(trial)
            hidden = generator.permutation(size)
            shuffled = text[np.ix_(hidden, hidden)]
            columns = grampa(visual, shuffled, 1.0)
            energy = ClosedFormEnergy(
                standardized(torch.from_numpy(shuffled).to(device)).float(),
                standardized(torch.from_numpy(visual).to(device)).float(),
                1.0, 1.0, 256,
            )
            final, _ = steepest_descent(
                energy, torch.from_numpy(columns.copy()).to(device), swaps, 3000
            )
            accuracies.append(
                float((hidden[final.cpu().numpy()] == np.arange(size)).mean())
            )
        row = {
            "rank": rank,
            "correlation": rho,
            "recovery": float(np.mean(accuracies)),
            "visual_effective_rank": effective_rank(visual),
        }
        rows.append(row)
        print(
            f"rank={rank:<5} rho={rho:.4f}  recovery={row['recovery']:.3f}  "
            f"eff_rank={row['visual_effective_rank']:.1f}",
            flush=True,
        )

    summary = {
        "protocol": (
            "Both fields truncated to the same rank, then matched blind from a "
            "spectral start with exact refinement. Hidden pairing generated "
            "per trial and read only for scoring."
        ),
        "label": args.label,
        "size": size,
        "chance": 1.0 / size,
        "rows": rows,
        "reading": (
            "Recovery collapsing at low rank while the correlation stays above "
            "the recovery threshold shows the headline correlation is not "
            "sufficient and the shared spectrum's width is the binding "
            "constraint."
        ),
    }
    write_json(args.output, summary)
    print(json.dumps({"chance": 1.0 / size}))


if __name__ == "__main__":
    main()