"""Natural-data field pipeline: continuous states and content projection.

Produces the field correlation reported for Visual Genome. Two choices
carry most of it, both measured. States stay continuous -- quantising
segments and phrases into class codes costs more than half the signal,
because relation fields need scene similarity and nothing else. And each
side is projected onto the directions that separate scenes rather than
parts, fitted per modality on scenes outside the evaluated set.

Vision states come from `natural_objects`; language states are
positive-PMI context vectors of the corpus, averaged per phrase. Hidden
pairs are read only to report the correlation.
"""

from __future__ import annotations

import argparse
import json
import re
from collections import Counter
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from scipy.linalg import eigh
from sklearn.decomposition import TruncatedSVD

from .common import read_json, seed_everything, write_json
from .natural_families import load_phrases, statistics
from .synth_cc_battery import moment_field


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vg-dir", default="artifacts/vg_5k")
    parser.add_argument("--objects", default="artifacts/vg_5k/natural_objects.pt")
    parser.add_argument("--samples", type=int, default=256)
    parser.add_argument("--fit-scenes", type=int, default=1200)
    parser.add_argument("--word-vectors", type=int, default=48)
    parser.add_argument("--min-count", type=int, default=60)
    parser.add_argument("--max-vocabulary", type=int, default=600)
    parser.add_argument("--context-window", type=int, default=2)
    parser.add_argument("--keep", type=int, default=64)
    parser.add_argument("--shrinkage", type=float, default=0.05)
    parser.add_argument(
        "--weight-power",
        type=float,
        default=0.5,
        help="Scale each kept direction by its discriminant eigenvalue to "
        "this power. 0 reproduces the unweighted projection.",
    )
    parser.add_argument("--views", type=int, default=1)
    parser.add_argument("--moment-degree", type=int, default=2, choices=[2, 3])
    parser.add_argument("--third-width", type=int, default=12)
    parser.add_argument(
        "--phrase-encoding", default="mean", choices=["mean", "structured"]
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--output", default="artifacts/vg_5k/natural_pipeline.pt")
    return parser.parse_args()


def word_vectors(
    phrases: list[list[str]], args: argparse.Namespace
) -> dict[str, np.ndarray]:
    """Positive-PMI context vectors, reduced by truncated SVD."""
    counts = Counter(token for tokens in phrases for token in set(tokens))
    vocabulary = [
        word
        for word, count in counts.most_common(args.max_vocabulary)
        if count >= args.min_count
    ]
    _, _, context = statistics(phrases, set(vocabulary), args.context_window)
    features = sorted({key for word in vocabulary for key in context[word]})
    index = {key: position for position, key in enumerate(features)}
    matrix = np.zeros((len(vocabulary), len(features)))
    for row, word in enumerate(vocabulary):
        for key, value in context[word].items():
            matrix[row, index[key]] = value
    total = matrix.sum()
    expected = matrix.sum(1, keepdims=True) * matrix.sum(0, keepdims=True) / total
    pmi = np.log(np.maximum(matrix, 1e-9) / np.maximum(expected, 1e-9))
    pmi[matrix == 0] = 0.0
    embedding = TruncatedSVD(args.word_vectors, random_state=args.seed).fit_transform(
        np.maximum(pmi, 0.0)
    )
    embedding /= np.linalg.norm(embedding, axis=1, keepdims=True).clip(1e-9)
    return {word: embedding[row] for row, word in enumerate(vocabulary)}


def moment_field_degree(
    sets: list[torch.Tensor], degree: int, third_width: int
) -> torch.Tensor:
    """Set kernel carrying moments up to `degree`.

    A kernel carrying third moments spans more directions than one carrying
    two, so raising the degree raises the field's rank by construction --
    which is the quantity the rank experiment showed to govern recovery. The
    third moment is taken over a narrower basis because its size is cubic;
    the first two are kept at full width.
    """
    if degree < 3:
        return moment_field(sets)
    phis = []
    for members in sets:
        first = members.mean(0)
        second = (members[:, :, None] * members[:, None, :]).mean(0).flatten()
        narrow = members[:, :third_width]
        outer = narrow[:, :, None] * narrow[:, None, :]
        third = (outer[:, :, :, None] * narrow[:, None, None, :]).mean(0).flatten()
        phi = torch.cat([first, second, third])
        phis.append(phi / phi.norm().clamp_min(1e-9))
    stacked = torch.stack(phis)
    return stacked @ stacked.T


def content_directions(
    sets: list[np.ndarray], shrinkage: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Directions separating scenes more than they separate parts.

    Returns the discriminant eigenvalues alongside the basis. Scaling each
    direction by its own eigenvalue lets a wide basis be kept without the
    weakly separating directions drowning the strong ones, which is what
    makes a wide keep beat a narrow one -- and a wide keep is what raises
    the rank of the resulting field.
    """
    means = np.stack([item.mean(0) for item in sets])
    centre = means.mean(0)
    within = np.concatenate([item - item.mean(0, keepdims=True) for item in sets])
    scatter_within = within.T @ within / max(len(within) - 1, 1)
    centred = means - centre
    scatter_between = centred.T @ centred / max(len(centred) - 1, 1)
    trace = np.trace(scatter_within) / len(scatter_within)
    values, vectors = eigh(
        scatter_between, scatter_within + shrinkage * trace * np.eye(len(scatter_within))
    )
    order = np.argsort(values)[::-1]
    return centre, vectors[:, order], np.maximum(values[order], 0.0)


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    vg_dir = Path(args.vg_dir)
    truth = [
        json.loads(line)
        for line in (vg_dir / "ground_truth.private.jsonl")
        .read_text(encoding="utf-8")
        .splitlines()
        if line.strip()
    ]
    records = {
        json.loads(line)["node_id"]: json.loads(line)
        for line in (vg_dir / "text_nodes.jsonl").read_text(encoding="utf-8").splitlines()
        if line.strip()
    }
    state = torch.load(args.objects, map_location="cpu", weights_only=False)
    per_view = state.get("view_segments") or [state["segments"]]
    views = min(args.views, len(per_view))
    index = {node: position for position, node in enumerate(state["node_ids"])}

    vectors = word_vectors(load_phrases(vg_dir / "text_nodes.jsonl", "region_closed"), args)

    width = args.word_vectors

    def encode_phrase(phrase: str) -> np.ndarray | None:
        """One vector per region description.

        Averaging every word of a phrase collapses "red bus" and "bus red"
        and, worse, mixes the thing named with what is said about it. The
        synthetic world never had to face this because its captions were
        encoded into separate factor blocks. Splitting the head from its
        modifiers recovers part of that structure from English word order
        alone -- the last token of an English noun phrase is its head -- which
        is a fact about one language, not a claim about the other modality,
        so it stays inside Tier 0.
        """
        tokens = [t for t in re.findall(r"[a-z]+", phrase.lower()) if t in vectors]
        if not tokens:
            return None
        if args.phrase_encoding == "mean":
            return np.mean([vectors[t] for t in tokens], axis=0)
        head = vectors[tokens[-1]]
        modifiers = (
            np.mean([vectors[t] for t in tokens[:-1]], axis=0)
            if len(tokens) > 1
            else np.zeros(width)
        )
        return np.concatenate([head, modifiers])

    def language_state(node: str) -> np.ndarray | None:
        rows = [
            row
            for row in (encode_phrase(p) for p in records[node]["region_closed"])
            if row is not None
        ]
        return np.stack(rows) if rows else None

    def vision_state(node: str, view: int) -> np.ndarray | None:
        segments = per_view[view][index[node]]
        return (
            np.stack([item["feature"] for item in segments]).astype(np.float64)
            if segments
            else None
        )

    pairs = [
        pair
        for pair in truth
        if pair["vision_node_id"] in index and pair["text_node_id"] in records
    ]
    fit = pairs[args.samples : args.samples + args.fit_scenes]
    language_fit = [language_state(p["text_node_id"]) for p in fit]
    language_fit = [item for item in language_fit if item is not None and len(item) > 1]
    vision_fit = [vision_state(p["vision_node_id"], 0) for p in fit]
    vision_fit = [item for item in vision_fit if item is not None and len(item) > 1]
    language_fitted = content_directions(language_fit, args.shrinkage)
    vision_fitted = content_directions(vision_fit, args.shrinkage)

    evaluated = [
        pair for pair in pairs[: args.samples]
        if language_state(pair["text_node_id"]) is not None
    ]
    keep = args.keep

    def project(raw: np.ndarray, fitted: tuple) -> torch.Tensor:
        centre, basis, values = fitted
        width = min(keep, basis.shape[1])
        scale = values[:width] ** args.weight_power if args.weight_power else 1.0
        return F.normalize(
            torch.tensor(
                (raw - centre) @ basis[:, :width] * scale, dtype=torch.float32
            ), dim=-1
        )

    text_sets = [
        project(language_state(p["text_node_id"]), language_fitted) for p in evaluated
    ]
    view_fields = []
    for view in range(views):
        sets = [
            project(vision_state(p["vision_node_id"], view), vision_fitted)
            for p in evaluated
        ]
        view_fields.append(
            moment_field_degree(sets, args.moment_degree, args.third_width)
        )
    visual_field = torch.stack(view_fields).mean(0)
    text_field = moment_field_degree(
        text_sets, args.moment_degree, args.third_width
    )

    mask = ~np.eye(len(evaluated), dtype=bool)
    correlation = float(
        np.corrcoef(
            visual_field.double().numpy()[mask], text_field.double().numpy()[mask]
        )[0, 1]
    )
    torch.save(
        {"visual_field": visual_field, "text_field": text_field,
         "nodes": [p["vision_node_id"] for p in evaluated]},
        args.output,
    )
    summary = {
        "protocol": (
            "Continuous states, content projection fitted per modality on "
            "scenes outside the evaluated set. Hidden pairs are read only "
            "for the correlation."
        ),
        "samples": len(evaluated),
        "views": views,
        "kept_directions": keep,
        "weight_power": args.weight_power,
        "moment_degree": args.moment_degree,
        "phrase_encoding": args.phrase_encoding,
        "field_correlation_at_truth": correlation,
        "recovery_threshold": 0.9,
    }
    print(json.dumps(summary))
    write_json(str(args.output).replace(".pt", ".json"), summary)


if __name__ == "__main__":
    main()