"""The other axis: recovery against correlation with the shared width intact.

The truncation ladder varied the width at nearly fixed correlation. This
varies the correlation at full width, by adding independent noise to both
fields, and the two together give the picture the single statistic could not.

The control matters because it guards against overcorrecting. Showing that
width is a separate necessary condition does not show that correlation stopped
mattering, and natural data is short on both -- correlation 0.725 against 0.93,
width 15 against 26. If full-width fields still recover at a correlation near
0.72, then width is what natural data lacks; if they do not, both gaps are real
and the correlation remains a target rather than a discredited one.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json
from .shared_rank import spectrum
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("--trials", type=int, default=3)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument(
        "--noise", type=float, nargs="+",
        default=[0.0, 0.2, 0.35, 0.5, 0.7, 0.9, 1.2],
    )
    parser.add_argument("--output", default="artifacts/synth_v1/rho_at_full_width.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 symmetric_noise(size: int, generator) -> np.ndarray:
    raw = generator.normal(size=(size, size))
    out = (raw + raw.T) / np.sqrt(2.0)
    np.fill_diagonal(out, 0.0)
    return out


def shared_count(visual: np.ndarray, text: np.ndarray, width: int = 32) -> int:
    _, visual_vectors = spectrum(visual)
    _, text_vectors = spectrum(text)
    cosines = np.clip(
        np.linalg.svd(
            visual_vectors[:, :width].T @ text_vectors[:, :width], compute_uv=False
        ), 0.0, 1.0,
    )
    return int((cosines > 0.7).sum())


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

    rows = []
    for level in args.noise:
        correlations, accuracies, widths = [], [], []
        for trial in range(args.trials):
            generator = np.random.default_rng(1000 + trial)
            visual = normalise(
                visual_clean + level * symmetric_noise(size, generator)
            )
            text = normalise(text_clean + level * symmetric_noise(size, generator))
            mask = ~np.eye(size, dtype=bool)
            correlations.append(float(np.corrcoef(visual[mask], text[mask])[0, 1]))
            widths.append(shared_count(visual, text))

            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, 4000
            )
            accuracies.append(
                float((hidden[final.cpu().numpy()] == np.arange(size)).mean())
            )
        row = {
            "noise": level,
            "correlation": float(np.mean(correlations)),
            "shared_directions": float(np.mean(widths)),
            "recovery": float(np.mean(accuracies)),
        }
        rows.append(row)
        print(
            f"noise={level:<5} rho={row['correlation']:.3f}  "
            f"shared={row['shared_directions']:5.1f}  "
            f"recovery={row['recovery']:.3f}",
            flush=True,
        )

    summary = {
        "protocol": (
            "Independent symmetric noise added to both fields, so the shared "
            "spectrum stays wide while the correlation falls. Hidden pairing "
            "generated per trial and read only for scoring."
        ),
        "size": size,
        "chance": 1.0 / size,
        "rows": rows,
    }
    print(json.dumps({"chance": 1.0 / size}))
    write_json(args.output, summary)


if __name__ == "__main__":
    main()