"""
Causal-patching / white-box localizer. Attempts M (mechanism) and S (structure).

WHITE-BOX: uses the encoder's ablate() and token_states() / encode() — the model's
own computation — never the held-out ground truth (U, angles, v_*).

M (localize where c is written): the concept field c is a SENTENCE-LEVEL additive
component. The technique discovers candidate concept directions unsupervised (ICA
on a corpus, like the D-axis), then ranks each ablatable layer by how much ablating
it COLLAPSES energy along those discovered concept directions. The layer whose
ablation collapses the concept component the most is named L_write. This mirrors
the causal answer key (concept-energy collapse profile) but is computed by the
technique from its OWN discovered directions + the public ablate() — no peeking.

S (recover R): probe the position operator by encoding single-token sentences at
positions 0 and 1. For a length-1 sentence z(t,p) = R^p E(t) + c(∅) = R^p E(t)
(no concepts). So R maps z(t,0) -> z(t,1) linearly across tokens: solve the least
-squares R_hat with z(·,1) ≈ R_hat z(·,0) over many tokens. This recovers the
authored rotation up to the per-plane sign symmetry — a joint-intervention recovery
of the generative subgraph.
"""

from __future__ import annotations

import numpy as np
from sklearn.decomposition import FastICA

from .base import Technique


class CausalPatch(Technique):
    NAME = "causal_patch"
    AXES = "MS"

    # ---- M: localize the write site ----
    def localize(self, encoder, items):
        # 1) discover candidate concept directions unsupervised on these items.
        Z = encoder.encode(items)
        Zc = Z - Z.mean(0, keepdims=True)
        k = 16
        try:
            ica = FastICA(n_components=k, random_state=0, max_iter=1000, tol=1e-3,
                          whiten="unit-variance").fit(Zc)
            comps = ica.mixing_.T
        except Exception:
            from sklearn.decomposition import PCA
            comps = PCA(n_components=k, random_state=0).fit(Zc).components_
        comps = comps / (np.linalg.norm(comps, axis=1, keepdims=True) + 1e-9)

        base_concept = float(np.mean((Z @ comps.T) ** 2)) + 1e-12
        base_total = float(np.mean(Z ** 2)) + 1e-12
        # 2) rank layers by SPECIFIC concept-energy collapse (concept frac removed
        #    minus total frac removed) so the trivial 'pool' sink that zeros
        #    everything is not mistaken for the write site.
        collapse = {}
        for layer in encoder.layers:
            Za = encoder.ablate(items, layer)
            concept_frac = 1.0 - float(np.mean((Za @ comps.T) ** 2)) / base_concept
            total_frac = 1.0 - float(np.mean(Za ** 2)) / base_total
            collapse[layer] = concept_frac - total_frac
        ranking = sorted(encoder.layers, key=lambda L: -collapse[L])
        # CALIBRATION-AWARE commit (v0.4): the write-site is real only if ONE layer
        # shows a clearly SPECIFIC concept collapse (positive specificity AND a clear
        # margin over the runner-up). For a DIFFUSE field (Cm phantom-site negative)
        # the profile is flat and the top specificity is ~0 or negative => no site =>
        # ABSTAIN, rather than naming a phantom. We measure the absolute specificity
        # margin between the best and second-best ABLATABLE layer (excluding 'pool',
        # the trivial sink that zeros everything and always has specificity ~0).
        ablatable = [L for L in ranking if L != "pool"]
        spec_sorted = sorted((collapse[L] for L in ablatable), reverse=True)
        top_spec = spec_sorted[0] if spec_sorted else 0.0
        margin = (spec_sorted[0] - spec_sorted[1]) if len(spec_sorted) > 1 else top_spec
        # a genuine write-site: positive specific collapse AND a real margin.
        has_site = (top_spec > 0.1) and (margin > 0.1)
        conf = float(np.clip(margin, 0, 1)) if has_site else 0.0
        return bool(has_site), {"ranking": ranking, "collapse": collapse}, conf

    # ---- S: recover the rotation operator ----
    def structure(self, encoder, data):
        K = data["K"]; d = data["d"]
        n_tok = min(data["n_probe"], K)
        rng = np.random.default_rng(0)
        toks = rng.choice(K, size=n_tok, replace=False)

        # single-token sentences at position 0 and at position 1.
        items0, items1 = [], []
        for t in toks:
            base = {"positions": [0], "tokens": [int(t)], "concepts": [],
                    "entities": {}, "number": None, "ro_g": 0.0, "causal_g": 0.0}
            items0.append({**base, "positions": [0]})
            items1.append({**base, "positions": [1]})
        Z0 = encoder.encode(items0)        # ~ R^0 E(t) = E(t)
        Z1 = encoder.encode(items1)        # ~ R^1 E(t)

        # solve Z1 ≈ Z0 @ R_hat^T  (row vectors): R_hat = (Z0^+ Z1)^T
        # least squares: R_hat^T = pinv(Z0) Z1
        RhatT, *_ = np.linalg.lstsq(Z0, Z1, rcond=None)   # (d, d): Z1 ≈ Z0 @ RhatT
        R_hat = RhatT.T
        # residual-based confidence
        resid = np.linalg.norm(Z0 @ RhatT - Z1) / (np.linalg.norm(Z1) + 1e-9)
        conf = float(np.clip(1.0 - resid, 0, 1))
        return True, {"R_hat": R_hat.tolist()}, conf