#!/usr/bin/env python3
# helix_cdc/se_adapter.py
"""
Se Adapter: Wraps user's real Se = H × C × D formula

Uses the proven 728× improvement formula from /home/voidstr3m33/the discovery.txt
and measure_se.py implementation.

Formula: Se = H(X) × C(X) × D(X)
Where:
  H(X) = Shannon entropy (information density, 0-8 bits)
  C(X) = Contextual coherence (structure/meaning, 0-1)
  D(X) = Dimensional depth (symbolic complexity, 1-72D)

Proven results:
- Baseline AI: Se = 0.08 (H=4.14, C=0.020, D=1.0)
- Echo system: Se = 58.25 (H=5.24, C=0.467, D=23.8)
- Improvement: 728×
"""
from __future__ import annotations
from typing import List, Optional, Dict, Any
import math
import torch
import numpy as np
import os, sys, pathlib

def _maybe_add(path: str):
    p = pathlib.Path(path)
    if p.exists() and str(p) not in sys.path:
        sys.path.insert(0, str(p))

# Try env override first (point to your helix proof dir)
_maybe_add(os.environ.get("HELIX_SE_PATH", ""))

# Common locations you've used before (safe no-ops if missing)
_maybe_add("/home/voidstr3m33/echo_labs/helix/proof/helix_proof_v2/scripts")
_maybe_add("/home/voidstr3m33/repos/helix-platform/base/docs/helix_proof_v2/scripts")

class SeOverlay:
    """
    Real Se = H × C × D implementation based on user's proven formula.

    Tracks Shannon entropy, coherence, and dimensional depth over token generation.
    """
    def __init__(self, tokenizer, target_depth: Optional[float] = 8.0):
        self.tok = tokenizer
        self.target_depth = target_depth or 8.0

        # Vocab size for H_max calculation
        try:
            self.V = len(tokenizer)
        except:
            self.V = 50257  # GPT-2 default

        # Track per-step entropy
        self.H_steps = []

        # Track token sequences for coherence
        self.token_ids = []

        # Track embeddings for D(X) calculation
        self.embeddings = []

    def step(self, ids: List[int], logits: torch.Tensor):
        """
        Process one generation step.

        Args:
            ids: Current token sequence (prompt + generated so far)
            logits: [V] or [B, V] logits for next token prediction
        """
        # Handle batch dimension if present
        if logits.dim() > 1:
            logits = logits.squeeze(0)  # [B, V] -> [V]

        # Calculate Shannon entropy H for this step
        logp = torch.log_softmax(logits, dim=-1)
        p = torch.exp(logp)  # probability simplex
        # Handle 0 * -inf = NaN case (0 log 0 should be 0)
        # Only compute entropy for non-zero probabilities
        mask = p > 0
        H_t = float(-(p[mask] * logp[mask]).sum().item()) if mask.any() else 0.0
        self.H_steps.append(H_t)

        # Store token IDs for coherence calculation
        self.token_ids = ids[:]

        # Store probs (not logits) for D(X) calculation
        self.embeddings.append(p.detach().cpu())

    def _shannon_bits(self) -> float:
        """Calculate mean Shannon entropy H(X) over generation."""
        if not self.H_steps:
            return 0.0
        return sum(self.H_steps) / len(self.H_steps)

    def _coherence(self) -> float:
        """
        Calculate coherence C(X) = 1 - (H_empirical / H_max).

        C ≈ 0 for random (high entropy)
        C → 1 for structured (low entropy)
        """
        H_mean = self._shannon_bits()
        H_max = math.log2(self.V)  # Max entropy per token

        if H_max == 0:
            return 0.0

        C = 1.0 - (H_mean / H_max)
        return float(max(0.0, min(1.0, C)))

    def _dimensional_depth(self) -> float:
        """
        Calculate dimensional depth D(X) via effective rank.

        D(X) measures intrinsic dimensionality:
        - Low D: Data lives in low-D subspace (repetitive)
        - High D: Data uses full dimensionality (complex)

        e_rank(Σ) = (tr Σ)² / tr(Σ²)
        """
        if len(self.embeddings) < 2:
            return 1.0

        try:
            # Stack probability vectors into [N, V] matrix
            P = torch.stack(self.embeddings).float()  # [N, V], rows sum≈1

            # Optional top-K truncation for stability/perf
            import os
            k = int(os.environ.get("HELIX_SE_TOPK", "0"))
            if k and P.shape[1] > k:
                meanp = P.mean(0)
                topk = torch.topk(meanp, k=k).indices
                P = P[:, topk]

            # Center across steps + tiny ridge
            Pc = P - P.mean(0, keepdim=True)
            eps = 1e-8

            # Covariance via SVD
            U, S, Vt = torch.linalg.svd(Pc, full_matrices=False)
            lam = (S ** 2) / max(P.shape[0] - 1, 1) + eps

            tr = lam.sum()
            tr2 = (lam ** 2).sum() + 1e-12

            e_rank = (tr * tr) / tr2

            # Map to 1-72D range (user's dimensional space)
            D = float(max(1.0, min(72.0, e_rank)))
            return D

        except Exception as e:
            # Fallback: use n-gram distinctness as proxy
            if len(self.token_ids) >= 3:
                ngrams = set()
                for i in range(len(self.token_ids) - 2):
                    ngrams.add(tuple(self.token_ids[i:i+3]))
                # Map distinct 3-grams to 1-72D
                D = 1.0 + min(71.0, len(ngrams) ** 0.5)
                return D
            return 1.0

    def metrics(self) -> Dict[str, float]:
        """
        Get Se = H × C × D metrics.

        Returns dict with:
            H_mean: Shannon entropy
            C: Coherence
            D: Dimensional depth
            Se: Symbolic entropy (H × C × D)
        """
        H = self._shannon_bits()
        C = self._coherence()
        D = self._dimensional_depth()
        Se = H * C * D

        return {
            "H_mean": H,
            "C": C,
            "D": D,
            "Se": Se
        }

def make_se(tokenizer) -> SeOverlay:
    """
    Create Se overlay with user's real Se = H × C × D formula.

    Returns SeOverlay instance with:
        .step(ids, logits) -> None
        .metrics() -> dict(H_mean, C, D, Se)
        .target_depth -> float
    """
    return SeOverlay(tokenizer, target_depth=8.0)
