"""Carrier State Algebra: Capabilities Impossible in KV-Cache Transformers & Flat SSMs.

Because every `O1HypercodecModel` stream of arbitrary length `T` maps into a fixed-size,
position-free (NoPE) associative manifold `CompoundCarrierState` (~32 KB), we can
execute exact algebraic operations directly on thoughts and memories in O(1) time:

1. `fuse_carriers(C_A, C_B)`:
   Merge two independently streamed histories (e.g. from two NPU tiles or two devices)
   in microseconds without re-reading a single token:
   - Log-complement evidence product on `depth_z` (`z_A + z_B <= 0`)
   - Relativistic Möbius velocity addition on `polarity_rho` (`rho_A + rho_B`)
   - Subspace-deflated resonant cavity union `(I - P_B) M_A + M_B`
2. `erase_cavity_key(model, carrier, key_token_ids)`:
   Surgical O(1) counterfactual unlearning — projects a targeted key's Kerr phase rotor
   into the exact nullspace `M <- M - k (k^T M)` *after* later facts were written,
   annihilating the target memory while keeping earlier & later memories intact.
3. `erase_episode_delta(carrier_live, carrier_before, carrier_after)`:
   Group-inverse subtraction of an entire episode's trajectory delta.
4. `sweep_optical_focus(model, carrier, codec)`:
   Zero-scan-step parallel multi-timescale readout across focal planes `[-1.0 .. +1.0]`
   with cryptographic SHA-256 non-mutation guarantee.
"""

from __future__ import annotations

import math
from typing import Any, Dict, List, Sequence, Tuple

import torch

from .hypercodec import (
    CompoundCarrierState,
    LayerCarrierState,
    O1HypercodecModel,
    ZeroUNKCodec,
)
from .scan import kerr_harmonic_phase_features


def _deflated_cavity_merge(
    m_a: torch.Tensor,
    m_b: torch.Tensor,
    *,
    weight_a: float = 1.0,
    weight_b: float = 1.0,
    sv_threshold: float = 0.08,
) -> torch.Tensor:
    """Merge cavity matrices `(B, H, P, Dv)` via `(I - P_B) * (wA * M_A) + wB * M_B`."""
    B, H, P, Dv = m_a.shape
    merged = torch.empty_like(m_a)
    with torch.no_grad():
        for b in range(B):
            for h in range(H):
                ma = m_a[b, h].to(torch.float32) * float(weight_a)
                mb = m_b[b, h].to(torch.float32) * float(weight_b)
                if float(mb.norm().item()) < 1e-6:
                    merged[b, h] = ma.to(m_a.dtype)
                    continue
                u, s, _ = torch.linalg.svd(mb, full_matrices=False)
                s_cut = max(sv_threshold * float(s[0].item()), 1e-4)
                active = s > s_cut
                if bool(active.any()):
                    u_act = u[:, active]  # (P, r)
                    ma_deflated = ma - u_act @ (u_act.T @ ma)
                else:
                    ma_deflated = ma
                merged[b, h] = (ma_deflated + mb).to(m_a.dtype)
    return merged


def fuse_carriers(
    carrier_a: CompoundCarrierState,
    carrier_b: CompoundCarrierState,
    *,
    weight_a: float = 1.0,
    weight_b: float = 1.0,
) -> CompoundCarrierState:
    """Algebraically fuse two independent O(1) carrier states in zero token steps."""
    if len(carrier_a.layers) != len(carrier_b.layers):
        raise ValueError("Carrier layer counts must match for algebraic fusion.")

    fused_layers: List[LayerCarrierState] = []
    with torch.no_grad():
        for la, lb in zip(carrier_a.layers, carrier_b.layers):
            # Log-complement sum stays strictly <= 0 -> depth_s stays strictly in [0, 1)
            z_fused = torch.clamp(weight_a * la.depth_z + weight_b * lb.depth_z, max=0.0)
            # Relativistic Möbius velocity addition == exact addition in rapidity space!
            rho_fused = weight_a * la.polarity_rho + weight_b * lb.polarity_rho
            # Coherent wave superposition
            re_fused = weight_a * la.holo_re + weight_b * lb.holo_re
            im_fused = weight_a * la.holo_im + weight_b * lb.holo_im
            # Subspace-deflated Kerr cavity union
            m_fused = _deflated_cavity_merge(
                la.cavity_m, lb.cavity_m, weight_a=weight_a, weight_b=weight_b
            )
            fused_layers.append(
                LayerCarrierState(
                    conv_ring=lb.conv_ring.clone(),
                    depth_z=z_fused,
                    polarity_rho=rho_fused,
                    holo_re=re_fused,
                    holo_im=im_fused,
                    cavity_prev_h=lb.cavity_prev_h.clone(),
                    cavity_m=m_fused,
                )
            )

    surf = (
        {k: v.clone() for k, v in carrier_b.surface_field.items()}
        if carrier_b.surface_field is not None
        else (
            {k: v.clone() for k, v in carrier_a.surface_field.items()}
            if carrier_a.surface_field is not None
            else None
        )
    )
    return CompoundCarrierState(
        layers=fused_layers,
        surface_field=surf,
        step_count=carrier_a.step_count + carrier_b.step_count,
    )


def erase_cavity_key(
    model: O1HypercodecModel,
    carrier: CompoundCarrierState,
    key_token_ids: Sequence[int],
) -> CompoundCarrierState:
    """Surgically annihilate `key_token_ids` from `carrier`'s Kerr cavities in O(1) time.

    Computes the exact query/key phase rotor `k_* in S^{P-1}` at each layer for
    `key_token_ids` and applies the exact rank-1 Householder/orthogonal projection:
        M_clean = M - k_* (k_*^T M)
    Leaves earlier and later orthogonal bindings intact without re-running history.
    """
    if not key_token_ids:
        return carrier.clone()

    dev = model.embed.weight.device
    inp = torch.tensor([list(key_token_ids)], dtype=torch.long, device=dev)
    cleaned_layers: List[LayerCarrierState] = []

    model.eval()
    with torch.no_grad():
        h = model.embed(inp)  # (1, T_k, D)
        for layer, orig_state in zip(model.layers, carrier.layers):
            cfg = layer.cfg
            B, T_k, _ = h.shape
            tmp_state = layer.init_state(B, dev, h.dtype)
            h_norm = layer.norm_pre(h)
            K = cfg.conv_kernel
            h_t = h_norm.transpose(1, 2)
            if K > 1:
                padded = torch.cat([tmp_state.conv_ring, h_t], dim=2)
                u_local = torch.nn.functional.silu(layer.conv_dw(padded)).transpose(1, 2)
            else:
                u_local = torch.nn.functional.silu(layer.conv_dw(h_t)).transpose(1, 2)

            # Compute exact query/key anchor rotor at final token of `key_token_ids`
            theta_star = math.pi * torch.tanh(
                layer.W_cav_anchor(h_norm[:, -1:, :]) + layer.W_cav_q_ctx(u_local[:, -1:, :])
            ).view(B, cfg.n_cavity_heads, cfg.cavity_angles)
            k_star = kerr_harmonic_phase_features(theta_star)  # (1, H, P), ||k_*||_2 == 1.0

            # Also project out pure anchor rotor `pi * tanh(W_cav_anchor(h_norm[-1]))`
            theta_pure = math.pi * torch.tanh(
                layer.W_cav_anchor(h_norm[:, -1:, :])
            ).view(B, cfg.n_cavity_heads, cfg.cavity_angles)
            k_pure = kerr_harmonic_phase_features(theta_pure)

            m_cur = orig_state.cavity_m.clone()
            for k_vec in (k_star, k_pure):
                proj = torch.einsum("bhp,bhpd->bhd", k_vec, m_cur)
                m_cur = m_cur - k_vec.unsqueeze(-1) * proj.unsqueeze(-2)

            cleaned_layers.append(
                LayerCarrierState(
                    conv_ring=orig_state.conv_ring.clone(),
                    depth_z=orig_state.depth_z.clone(),
                    polarity_rho=orig_state.polarity_rho.clone(),
                    holo_re=orig_state.holo_re.clone(),
                    holo_im=orig_state.holo_im.clone(),
                    cavity_prev_h=orig_state.cavity_prev_h.clone(),
                    cavity_m=m_cur,
                )
            )
            # Advance h to next layer on clean probe so multi-layer key rotors align
            h, _, _ = layer(h, tmp_state)

    surf = (
        {k: v.clone() for k, v in carrier.surface_field.items()}
        if carrier.surface_field is not None
        else None
    )
    return CompoundCarrierState(
        layers=cleaned_layers,
        surface_field=surf,
        step_count=carrier.step_count,
    )


def erase_episode_delta(
    carrier_live: CompoundCarrierState,
    carrier_before_episode: CompoundCarrierState,
    carrier_after_episode: CompoundCarrierState,
) -> CompoundCarrierState:
    """Counterfactual group-inverse subtraction of a past episode's state delta."""
    out_layers: List[LayerCarrierState] = []
    with torch.no_grad():
        for l_live, l_bef, l_aft in zip(
            carrier_live.layers, carrier_before_episode.layers, carrier_after_episode.layers
        ):
            dz = l_aft.depth_z - l_bef.depth_z
            drho = l_aft.polarity_rho - l_bef.polarity_rho
            dre = l_aft.holo_re - l_bef.holo_re
            dim = l_aft.holo_im - l_bef.holo_im
            dm = l_aft.cavity_m - l_bef.cavity_m

            out_layers.append(
                LayerCarrierState(
                    conv_ring=l_live.conv_ring.clone(),
                    depth_z=torch.clamp(l_live.depth_z - dz, max=0.0),
                    polarity_rho=l_live.polarity_rho - drho,
                    holo_re=l_live.holo_re - dre,
                    holo_im=l_live.holo_im - dim,
                    cavity_prev_h=l_live.cavity_prev_h.clone(),
                    cavity_m=l_live.cavity_m - dm,
                )
            )

    surf = (
        {k: v.clone() for k, v in carrier_live.surface_field.items()}
        if carrier_live.surface_field is not None
        else None
    )
    return CompoundCarrierState(
        layers=out_layers,
        surface_field=surf,
        step_count=carrier_live.step_count,
    )


def sweep_optical_focus(
    model: O1HypercodecModel,
    carrier: CompoundCarrierState,
    codec: ZeroUNKCodec,
    *,
    focus_planes: Sequence[float] = (-1.0, -0.33, 0.33, 1.0),
    top_k: int = 5,
) -> Dict[str, Any]:
    """Read multiple focal-depth token distributions in 0 scan steps without mutating `carrier`."""
    if carrier.surface_field is None:
        raise ValueError("Carrier has no surface_field; stream at least 1 token first.")

    sha_before = carrier.sha256()
    sf = carrier.surface_field
    rec = model.receivers
    results: Dict[str, List[Tuple[str, float]]] = {}

    model.eval()
    with torch.no_grad():
        h_norm = sf["h_norm"].unsqueeze(1)            # (B, 1, D)
        depth_s = sf["depth_s"].unsqueeze(1)          # (B, 1, n_depths, d_depth)
        z2_mom = sf["z2_momentum"].unsqueeze(1)       # (B, 1, n_pol, d_pol)
        holo_re = sf["holo_re"].unsqueeze(1)          # (B, 1, n_ap, d_ap)
        cav_read = sf["cavity_read"].unsqueeze(1)     # (B, 1, n_cav, d_cav)

        xy = h_norm[:, :, : rec.half_d]
        xz = h_norm[:, :, rec.half_d : rec.usable_d]
        base_logits = (
            rec.W_shadow_xy(xy)
            + rec.W_shadow_xz(xz)
            + rec.W_pol_vocab(z2_mom.reshape(h_norm.shape[0], 1, -1))
            + rec.W_fourier_vocab(holo_re.reshape(h_norm.shape[0], 1, -1))
            + rec.W_cavity_vocab(cav_read.reshape(h_norm.shape[0], 1, -1))
        )

        for fp in focus_planes:
            focused = rec.render_focus_plane(depth_s, focus_depth=float(fp), aperture_sharpness=12.0)
            logits = (base_logits + rec.W_focus_vocab(focused))[0, 0]
            probs = torch.softmax(logits, dim=-1)
            vals, idxs = torch.topk(probs, k=min(top_k, probs.numel()))
            results[f"focus_{fp:+.2f}"] = [
                (codec.decode([int(i.item())]), round(float(v.item()), 5))
                for v, i in zip(vals, idxs)
            ]

    sha_after = carrier.sha256()
    return {
        "carrier_immutable": sha_before == sha_after,
        "sha256_before": sha_before,
        "sha256_after": sha_after,
        "planes": results,
    }