"""
holo-invariance
===============
Measure the invariance group of any model.

A mode is characterized by what it can ignore, not by what it can do.
This tool measures the invariance group directly: the set of input
transformations under which a model's operation is unchanged.

Three things the tool measures, for each transformation:

    mean_similarity   Continuous. Average closeness of outputs.
    invariance        Fraction of trials with similarity above a
                      threshold. This defines group membership.
    decision_change   Fraction of trials where the discrete output
                      changed. The "accuracy drop" analog.

The three can move independently. A model can be invariant in operation
but not in decisions (it operates the same way, its argmax flips). Or
invariant in decisions but not operation (different paths, same answer).
The signature of a model is its invariance vector across the
transformation set.

Fixes over the previous version:

  1. Signature similarity uses centered vectors. Prevents the constant
     dimensions from dominating the metric.

  2. Effective trial counting. Each transformation reports how often it
     actually changed the input. A transformation that did nothing on the
     test set reports invariance = 1.0 with n_effective = 0, and the
     verdict excludes it from group membership.

  3. Multiple comparators. Cosine alone misses ratio changes in
     normalized vectors. The report includes cosine, relative L1, and
     discrete decision change. The user picks which defines invariance.

Usage:

    from holo_invariance import InvarianceProbe, TEXT_TRANSFORMS

    def my_model(text):
        return np.array([...])   # vector, scalar, or string

    probe = InvarianceProbe(my_model, TEXT_TRANSFORMS)
    report = probe.run(my_texts)
    report.print()
    report.save_json("report.json")
    report.plot("signature.png")

For LLMs, wrap them:

    from holo_invariance import wrap_llm

    model = wrap_llm(
        api_call=lambda t: my_llm_call(t),
        feature_fn=lambda r: my_embedding(r),
    )

CLI:

    python holo_invariance.py                 # run demo
    python holo_invariance.py --help

Requires: numpy, matplotlib.
"""

from __future__ import annotations

import argparse
import json
import os
import string
from dataclasses import dataclass, field, asdict
from typing import Any, Callable, Dict, List, Optional

import numpy as np


# ============================================================
# Comparators — determine what "same output" means
# ============================================================
def cosine_similarity(a, b) -> float:
    """Cosine of two vectors. Range [-1, 1]."""
    a = np.asarray(a, dtype=float).ravel()
    b = np.asarray(b, dtype=float).ravel()
    if a.shape != b.shape:
        return 0.0
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na < 1e-12 or nb < 1e-12:
        return 0.0
    return float(a @ b / (na * nb))


def relative_l1(a, b) -> float:
    """Relative L1 distance, mapped to similarity in [0, 1]."""
    a = np.asarray(a, dtype=float).ravel()
    b = np.asarray(b, dtype=float).ravel()
    if a.shape != b.shape:
        return 0.0
    denom = np.linalg.norm(a) + np.linalg.norm(b) + 1e-12
    l1 = np.abs(a - b).sum()
    return float(1.0 - l1 / denom)


def scalar_similarity(a, b) -> float:
    a = float(a)
    b = float(b)
    denom = abs(a) + abs(b) + 1e-12
    return float(1.0 - abs(a - b) / denom)


def exact_match(a, b) -> float:
    return 1.0 if a == b else 0.0


def auto_compare(a, b) -> float:
    """Dispatch on type and shape."""
    if isinstance(a, str) and isinstance(b, str):
        return exact_match(a, b)
    try:
        aa = np.asarray(a)
        bb = np.asarray(b)
        if aa.ndim == 0 and bb.ndim == 0:
            return scalar_similarity(aa.item(), bb.item())
        if aa.shape == bb.shape and aa.size > 1:
            return cosine_similarity(aa.ravel(), bb.ravel())
    except Exception:
        pass
    return exact_match(a, b)


def discrete_readout(x, scalar_tol: float = 0.10):
    """
    Extract a discrete decision from any output.

    Vectors: argmax index.
    Scalars: relative-magnitude bin, so small perturbations do not
             register as decision changes.
    Strings: the string itself.
    """
    if isinstance(x, str):
        return x
    arr = np.asarray(x)
    if arr.ndim == 0:
        v = float(arr.item())
        if abs(v) < 1e-9:
            return 0
        return int(np.floor(np.log(abs(v) + 1e-9) / np.log(1 + scalar_tol)))
    if arr.size > 1:
        return int(np.argmax(arr))
    return arr.item()


# ============================================================
# Text transformations
# ============================================================
STOPWORDS = frozenset("""a an the of in on at to for and or but if so as is
are was were be been being has have had do does did this that these those
from with by about into over under""".split())

SYNONYMS = {
    "happy": "glad", "sad": "unhappy", "big": "large", "small": "tiny",
    "fast": "quick", "slow": "sluggish", "good": "fine", "bad": "poor",
    "buy": "purchase", "sell": "vend", "start": "begin", "stop": "cease",
    "make": "create", "use": "utilize", "get": "obtain", "give": "provide",
    "problem": "issue", "answer": "response", "question": "query",
    "show": "display", "hide": "conceal", "help": "assist", "find": "locate",
    "man": "gentleman", "woman": "lady", "child": "kid", "house": "home",
    "car": "auto", "dog": "canine", "cat": "feline", "food": "nourishment",
    "joy": "delight", "sorrow": "grief", "delight": "joy", "glad": "happy",
}


def t_case_flip(x, rng):
    if isinstance(x, str):
        return x.upper() if rng.random() < 0.5 else x.lower()
    return x

def t_whitespace(x, rng):
    if isinstance(x, str):
        sep = "  " if rng.random() < 0.5 else "\t"
        return sep.join(x.split())
    return x

def t_stopword_drop(x, rng):
    if isinstance(x, str):
        words = [w for w in x.split() if w.lower() not in STOPWORDS]
        return " ".join(words) if words else x
    return x

def t_punct_strip(x, rng):
    if isinstance(x, str):
        return x.translate(str.maketrans("", "", string.punctuation))
    return x

def t_punct_add(x, rng):
    """Add punctuation to exercise punctuation-invariance in either
    direction."""
    if isinstance(x, str):
        parts = x.split()
        if not parts:
            return x
        parts[0] = parts[0] + ","
        if len(parts) > 1:
            parts[-1] = parts[-1] + "."
        return " ".join(parts)
    return x

def t_synonym_swap(x, rng):
    if isinstance(x, str):
        out = []
        for w in x.split():
            stripped = w.strip(string.punctuation)
            lw = stripped.lower()
            if lw in SYNONYMS and rng.random() < 0.8:
                rep = SYNONYMS[lw]
                if w[:1].isupper():
                    rep = rep.capitalize()
                out.append(rep)
            else:
                out.append(w)
        return " ".join(out)
    return x

def t_word_shuffle(x, rng):
    if isinstance(x, str):
        words = x.split()
        if len(words) < 2:
            return x
        perm = rng.permutation(len(words))
        return " ".join(words[i] for i in perm)
    return x

def t_char_typo(x, rng, rate: float = 0.05):
    if isinstance(x, str):
        chars = list(x)
        for i in range(len(chars)):
            if rng.random() < rate and chars[i].isalpha():
                chars[i] = chr(ord('a') + int(rng.integers(0, 26)))
        return "".join(chars)
    return x

def t_length_double(x, rng):
    if isinstance(x, str):
        return x + " " + x
    return x

def t_case_flip_letters(x, rng):
    """Random per-letter case flip. Exercises case robustness more."""
    if isinstance(x, str):
        out = []
        for c in x:
            if c.isalpha() and rng.random() < 0.5:
                out.append(c.upper() if c.islower() else c.lower())
            else:
                out.append(c)
        return "".join(out)
    return x


TEXT_TRANSFORMS = {
    "case_flip":       t_case_flip,
    "case_flip_letters": t_case_flip_letters,
    "whitespace":      t_whitespace,
    "stopword_drop":   t_stopword_drop,
    "punct_strip":     t_punct_strip,
    "punct_add":       t_punct_add,
    "synonym_swap":    t_synonym_swap,
    "word_shuffle":    t_word_shuffle,
    "char_typo_5pct":  t_char_typo,
    "length_double":   t_length_double,
}


# ============================================================
# Numeric transformations
# ============================================================
def n_permute(x, rng):
    if isinstance(x, np.ndarray) and x.ndim == 1 and len(x) > 1:
        return x[rng.permutation(len(x))]
    return x

def n_scale(x, rng):
    if isinstance(x, np.ndarray):
        return x * 1.5
    return x

def n_shift(x, rng):
    if isinstance(x, np.ndarray):
        return x + 0.5
    return x

def n_sign_flip(x, rng):
    if isinstance(x, np.ndarray):
        return -x
    return x

def n_noise(x, rng, level: float = 0.1):
    if isinstance(x, np.ndarray):
        return x + level * rng.standard_normal(x.shape)
    return x

def n_zero_pad(x, rng):
    if isinstance(x, np.ndarray) and x.ndim == 1:
        return np.concatenate([x, np.zeros(1)])
    return x


NUMERIC_TRANSFORMS = {
    "permute":     n_permute,
    "scale_1.5":   n_scale,
    "shift_+0.5":  n_shift,
    "sign_flip":   n_sign_flip,
    "noise_0.1":   n_noise,
    "zero_pad":    n_zero_pad,
}


# ============================================================
# Probe
# ============================================================
@dataclass
class TransformResult:
    name: str
    n_trials: int
    n_effective: int
    mean_similarity: float
    invariance: float
    decision_change: float
    in_group: bool
    note: str = ""


@dataclass
class ProbeReport:
    model_name: str
    sim_threshold: float
    transform_names: List[str]
    results: Dict[str, TransformResult] = field(default_factory=dict)

    def signature(self) -> np.ndarray:
        return np.array([self.results[t].invariance
                        for t in self.transform_names])

    def effective_signature(self) -> np.ndarray:
        return np.array([
            self.results[t].invariance if self.results[t].n_effective > 0
            else np.nan
            for t in self.transform_names
        ])

    def print(self, title: Optional[str] = None):
        title = title or f"Invariance group: {self.model_name}"
        print(f"\n  {title}")
        print("  " + "-" * 88)
        print(f"  {'transformation':>20}  {'similarity':>10}  "
              f"{'invariance':>10}  {'Δdecision':>10}  "
              f"{'n_eff':>6}  {'in group':>10}")
        print("  " + "-" * 88)
        ordered = sorted(self.results.values(),
                         key=lambda r: (-r.in_group, -r.invariance))
        for r in ordered:
            verdict = "YES" if r.in_group else "no"
            if r.n_effective == 0:
                verdict = "n/a"
            print(f"  {r.name:>20}  {r.mean_similarity:>10.3f}  "
                  f"{r.invariance:>10.3f}  {r.decision_change:>10.3f}  "
                  f"{r.n_effective:>6}  {verdict:>10}")

    def to_dict(self) -> Dict:
        return {
            "model_name": self.model_name,
            "sim_threshold": self.sim_threshold,
            "signature": self.signature().tolist(),
            "effective_signature": [
                None if np.isnan(v) else float(v)
                for v in self.effective_signature()
            ],
            "transform_names": list(self.transform_names),
            "results": {k: asdict(v) for k, v in self.results.items()},
        }

    def save_json(self, path: str):
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)
        print(f"  saved: {path}")

    def plot(self, path: str):
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        names = [r.name for r in self.results.values()]
        invs = [r.invariance for r in self.results.values()]
        eff = [r.n_effective for r in self.results.values()]
        order = np.argsort(invs)
        names = [names[i] for i in order]
        invs = [invs[i] for i in order]
        eff = [eff[i] for i in order]
        colors = ["#2a9d3f" if (v >= 0.75 and e > 0) else
                  ("#888888" if e == 0 else "#c44e4e")
                  for v, e in zip(invs, eff)]
        fig, ax = plt.subplots(figsize=(10, max(3, 0.35 * len(names))))
        y = np.arange(len(names))
        ax.barh(y, invs, color=colors, edgecolor="black", lw=0.6)
        ax.axvline(self.sim_threshold, color="black", ls="--", lw=1,
                   label=f"threshold = {self.sim_threshold:.2f}")
        ax.set_yticks(y)
        ax.set_yticklabels(names)
        ax.invert_yaxis()
        ax.set_xlabel("invariance")
        ax.set_xlim(0, 1.05)
        ax.set_title(f"Invariance group: {self.model_name}\n"
                     f"green = in group, red = out, gray = transform no-op")
        ax.legend(loc="lower right", fontsize=8)
        ax.grid(alpha=0.3, axis="x")
        plt.tight_layout()
        plt.savefig(path, dpi=110, bbox_inches="tight")
        plt.close()
        print(f"  plot: {path}")


class InvarianceProbe:
    """
    Measure the invariance group of a model.

    Parameters
    ----------
    model : Callable[[input], output]
        Any model. Output can be a vector, a scalar, or a string.
    transforms : dict
        {name: Callable[[input, rng], input]}
    comparator : Callable[[output, output], float]
        Similarity in [0, 1] between two outputs. Default: auto_compare.
    sim_threshold : float
        Similarity above which a trial counts as invariant.
    seed : int
        RNG seed for reproducibility.
    """

    def __init__(self, model: Callable, transforms: Dict,
                 comparator: Callable = auto_compare,
                 sim_threshold: float = 0.90,
                 seed: int = 0,
                 name: Optional[str] = None):
        self.model = model
        self.transforms = transforms
        self.comparator = comparator
        self.sim_threshold = sim_threshold
        self.rng = np.random.default_rng(seed)
        self.name = name or getattr(model, "__class__", type(model)).__name__

    def _transform_changed_input(self, x, x_prime) -> bool:
        """Did the transformation actually change the input?"""
        if x is x_prime:
            return False
        if isinstance(x, str) and isinstance(x_prime, str):
            return x != x_prime
        try:
            a = np.asarray(x)
            b = np.asarray(x_prime)
            if a.shape != b.shape:
                return True
            return not np.allclose(a, b)
        except Exception:
            return x != x_prime

    def run(self, inputs: List[Any], n_repeats: int = 3) -> ProbeReport:
        transform_names = list(self.transforms.keys())
        raw = {name: {"sims": [], "dc": [], "changed": 0}
               for name in transform_names}

        for x in inputs:
            reference = self.model(x)
            ref_readout = discrete_readout(reference)
            for name, tr in self.transforms.items():
                for _ in range(n_repeats):
                    try:
                        x_prime = tr(x, self.rng)
                        if self._transform_changed_input(x, x_prime):
                            raw[name]["changed"] += 1
                        out = self.model(x_prime)
                        sim = self.comparator(reference, out)
                        dc = 0.0 if discrete_readout(out) == ref_readout else 1.0
                    except Exception:
                        sim, dc = 0.0, 1.0
                    raw[name]["sims"].append(sim)
                    raw[name]["dc"].append(dc)

        report = ProbeReport(
            model_name=self.name,
            sim_threshold=self.sim_threshold,
            transform_names=transform_names,
        )
        for name, r in raw.items():
            sims = np.array(r["sims"])
            dcs = np.array(r["dc"])
            n_trials = len(sims)
            n_eff = r["changed"]
            invariance = float((sims >= self.sim_threshold).mean())
            if n_eff == 0:
                in_group = False
                note = "transform no-op on input set"
            else:
                in_group = invariance >= 0.75
                note = ""
            report.results[name] = TransformResult(
                name=name,
                n_trials=n_trials,
                n_effective=n_eff,
                mean_similarity=float(sims.mean()),
                invariance=invariance,
                decision_change=float(dcs.mean()),
                in_group=in_group,
                note=note,
            )
        return report


# ============================================================
# Multi-model comparison
# ============================================================
def centered_signature_similarity(sig_a: np.ndarray,
                                   sig_b: np.ndarray,
                                   sig_all: np.ndarray) -> float:
    """
    Centered cosine similarity between two signature vectors.
    Subtracts the mean across models in each dimension so the
    difference lives where the models actually vary.
    """
    # sig_all: shape (n_models, n_transforms)
    mean_per_transform = np.nanmean(sig_all, axis=0)
    a = sig_a - mean_per_transform
    b = sig_b - mean_per_transform
    a = np.nan_to_num(a, nan=0.0)
    b = np.nan_to_num(b, nan=0.0)
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na < 1e-9 or nb < 1e-9:
        return 0.0
    return float(a @ b / (na * nb))


def compare_models(reports: Dict[str, ProbeReport]) -> None:
    """Print a side-by-side signature comparison of multiple reports."""
    names = list(reports.keys())
    transform_names = reports[names[0]].transform_names

    print("\n  " + "=" * 88)
    print("  MODEL SIGNATURES")
    print("  " + "=" * 88)
    print()
    print("  " + f"{'model':>16}  " +
          "  ".join(f"{t[:9]:>9}" for t in transform_names))
    sigs = []
    for n in names:
        sig = reports[n].effective_signature()
        sigs.append(sig)
        row = [("   n/a  " if np.isnan(v) else f"{v:>9.2f}") for v in sig]
        print(f"  {n:>16}  " + "  ".join(row))

    sigs = np.array(sigs)
    print()
    print("  Centered signature similarity:")
    print("  " + f"{'':>16}  " +
          "  ".join(f"{n[:9]:>9}" for n in names))
    for i, a in enumerate(names):
        row = []
        for j, b in enumerate(names):
            s = centered_signature_similarity(sigs[i], sigs[j], sigs)
            row.append(f"{s:>9.2f}")
        print(f"  {a:>16}  " + "  ".join(row))


# ============================================================
# LLM wrapper
# ============================================================
def wrap_llm(api_call: Callable, feature_fn: Optional[Callable] = None):
    """
    Wrap a text-in / text-out model for use with InvarianceProbe.

    api_call   : Callable[[str], Any] -- the raw model call.
    feature_fn : optional Callable[[Any], np.ndarray] -- converts the
                 raw output into a vector for comparison.

    If feature_fn is None, the raw output is compared (exact string
    match by default). This is too strict for free-form text. Recommend
    providing a feature_fn that returns an embedding or a token
    distribution.
    """
    def wrapped(x):
        out = api_call(x)
        if feature_fn is not None:
            return feature_fn(out)
        return out
    return wrapped


# ============================================================
# Demo models — chosen so their invariance groups differ
# ============================================================
class ExactMatchClassifier:
    """Only matches literal keywords."""
    WORDS = {
        "positive": ["happy", "joy", "delight"],
        "negative": ["sad", "sorrow", "grief"],
        "fast":     ["fast", "quick", "rapid"],
        "slow":     ["slow", "sluggish"],
    }

    def __call__(self, text):
        t = text.lower()
        s = np.zeros(4)
        for i, label in enumerate(self.WORDS):
            for w in self.WORDS[label]:
                if w in t:
                    s[i] += 1
        s = s + 0.05
        return s / s.sum()


class ConceptClassifier:
    """Broad synonym sets. Synonym-invariant by construction."""
    CONCEPTS = {
        "positive": ["happy", "glad", "joy", "delight", "fine", "great"],
        "negative": ["sad", "unhappy", "sorrow", "grief", "poor"],
        "fast":     ["fast", "quick", "rapid", "speedy"],
        "slow":     ["slow", "sluggish", "leisurely"],
    }

    def __call__(self, text):
        t = text.lower()
        s = np.zeros(4)
        for i, label in enumerate(self.CONCEPTS):
            for w in self.CONCEPTS[label]:
                if w in t:
                    s[i] += 1
        s = s + 0.05
        return s / s.sum()


class LengthOnlyClassifier:
    """Uses text length as its only feature."""
    def __call__(self, text):
        n = len(text.split())
        s = np.array([n, n * n, np.sqrt(n), np.log(n + 1)])
        return s / s.sum()


class RandomModel:
    """Control. Ignores input entirely."""
    def __call__(self, text):
        rng = np.random.default_rng(hash(text) & 0xffffffff)
        return rng.random(4)


# ============================================================
# Demo input set — includes punctuation, varied length, mixed case
# ============================================================
TEXT_INPUTS = [
    "I am happy today.",
    "The dog is fast and happy!",
    "She feels sad about the slow car, but the day is fine.",
    "A glad child runs quick.",
    "The morning brings joy and delight!",
    "This is a good day for a fast ride.",
    "Unhappy people walk slow, and the grey cloud follows.",
    "Sorrow fills the large house; the rooms are quiet.",
    "Quick thinking leads to a fine result, they say.",
    "The sluggish cat sleeps a lot during the long afternoon.",
    "Delight and joy are everywhere in the small town.",
    "Rapid movement startles the child; she laughs, then cries.",
    "A slow sorrow spreads through the house, room by room.",
    "Happy people find delight in small things, even on grey days.",
    "The grief of loss is slow to heal, but it heals.",
]


# ============================================================
# CLI
# ============================================================
def run_demo(output_dir: str = "holo_invariance_out") -> None:
    os.makedirs(output_dir, exist_ok=True)
    print("=" * 88)
    print("HOLO-INVARIANCE")
    print("Measuring the invariance group of any model.")
    print("=" * 88)

    models = {
        "ExactMatch": ExactMatchClassifier(),
        "Concept":    ConceptClassifier(),
        "LengthOnly": LengthOnlyClassifier(),
        "Random":     RandomModel(),
    }

    reports = {}
    for name, model in models.items():
        print("\n" + "=" * 88)
        print(f"MODEL: {name}")
        print("=" * 88)
        probe = InvarianceProbe(model, TEXT_TRANSFORMS,
                                sim_threshold=0.90, seed=0, name=name)
        report = probe.run(TEXT_INPUTS, n_repeats=3)
        report.print()
        report.save_json(os.path.join(output_dir, f"{name}.json"))
        report.plot(os.path.join(output_dir, f"{name}.png"))
        reports[name] = report

    compare_models(reports)

    # Numeric demo
    print("\n" + "=" * 88)
    print("MODEL: f(x) = sum(x^2)")
    print("=" * 88)
    def numeric_fn(x):
        return float(np.sum(x ** 2))
    numeric_inputs = [np.random.default_rng(i).standard_normal(8)
                      for i in range(20)]
    probe = InvarianceProbe(numeric_fn, NUMERIC_TRANSFORMS,
                            sim_threshold=0.90, seed=0,
                            name="SumOfSquares")
    report = probe.run(numeric_inputs, n_repeats=3)
    report.print()
    report.save_json(os.path.join(output_dir, "SumOfSquares.json"))
    report.plot(os.path.join(output_dir, "SumOfSquares.png"))

    print("\n" + "=" * 88)
    print(f"All outputs saved to: {output_dir}/")
    print("=" * 88)


def main():
    parser = argparse.ArgumentParser(
        description="Measure the invariance group of any model.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python holo_invariance.py
    python holo_invariance.py --output results/
    python holo_invariance.py --threshold 0.85

Library usage:

    from holo_invariance import InvarianceProbe, TEXT_TRANSFORMS

    probe = InvarianceProbe(my_model, TEXT_TRANSFORMS)
    report = probe.run(my_texts)
    report.print()
    report.save_json("report.json")
        """)
    parser.add_argument("--output", default="holo_invariance_out",
                        help="output directory (default: holo_invariance_out)")
    parser.add_argument("--threshold", type=float, default=0.90,
                        help="similarity threshold for invariance "
                             "(default: 0.90)")
    parser.add_argument("--demo", action="store_true",
                        help="run demo (default behavior)")
    args = parser.parse_args()

    run_demo(output_dir=args.output)


if __name__ == "__main__":
    main()