#!/usr/bin/env python3
"""BMCTE v2: Phi-NUVMAP Sparse Rollup.

Core insight: eigensolid = crossStep(s) = s. The state space at p/N=1/7
is sparse under phi-shell projection. We save each step to NUVMAP, allowing
resumption from any point via sparse addresses.

For p=11429: we DON'T enumerate 2^p masks. Instead we sample random masks
and record the evolving partial sum to NUVMAP sparse addresses.
"""

import argparse
import json
import math
from pathlib import Path

import numpy as np

RECEIPT_PATH = Path(__file__).resolve().parent / "extension_v2_nuvmap_receipt.json"
NUVMAP_DIR = Path(__file__).resolve().parent / "nuvmap_sparse"

def phi_encode(step: int, shell_capacity: int = 1) -> int:
    """Phi-shell address encoding: linear level growth."""
    level = step // shell_capacity
    index = step % shell_capacity
    return level * level + index

def build_isometry(N: int, p: int, seed: int) -> np.ndarray:
    rng = np.random.RandomState(seed)
    A = rng.randn(N, p) + 1j * rng.randn(N, p)
    Q, R = np.linalg.qr(A)
    U = Q @ np.diag(np.exp(1j * np.angle(np.diag(R))))
    return U

def run_bmcte_sparse(N: int, p: int, K: int, seed: int) -> dict:
    """Sparse BMCTE at p/N=1/7 threshold.
    
    Records each step to NUVMAP via phi-shell encoding.
    """
    U = build_isometry(N, p, seed)
    NUVMAP_DIR.mkdir(exist_ok=True)
    
    rng = np.random.RandomState(seed + 1000)
    col_probs = [np.abs(U[:, j])**2 for j in range(p)]
    
    weights = []
    saved_states = 0
    
    for shot in range(K):
        S = [int(np.searchsorted(np.cumsum(col_probs[j]), rng.random())) for j in range(p)]
        M = U[np.array(S), :]
        
        # Sample 100 random masks for partial sum
        partial = complex(0.0)
        for i in range(min(100, K * 10)):
            mask = np.random.RandomState(shot * 100 + i).randint(1, 2**32)
            # Extract bits up to p
            cols = [j for j in range(p) if (mask >> j) & 1]
            if not cols:
                continue
            sign = (-1) ** (p - len(cols))
            row_terms = [sum(M[row, j] for j in cols) for row in range(p)]
            partial += sign * complex(np.prod(row_terms))
        
        w = float(abs(partial) ** 2)
        weights.append(w)
        
        # Save to NUVMAP via phi-shell encoding
        addr = phi_encode(shot, 10)
        state = {
            "nuvmap_addr": addr,
            "step": shot,
            "partial_real": partial.real,
            "partial_imag": partial.imag,
            "weight": w,
        }
        (NUVMAP_DIR / f"step_{shot}.json").write_text(json.dumps(state))
        saved_states += 1
    
    # Compute metrics
    weights = np.array(weights)
    total = max(weights.sum(), 1e-15)
    probs = np.clip(weights / total, 1e-15, 1.0)
    entropy = float(-np.sum(probs * np.log2(probs))) if total > 0 else 0.0
    
    lambda_theory = math.exp(-p * p / N)
    
    return {
        "N": N, "p": p, "K": K, "seed": seed,
        "lambda_theory": lambda_theory,
        "p_over_N": p / N,
        "entropy_measured": entropy,
        "nuvmap_states_saved": saved_states,
        "at_threshold": abs(p/N - 1/7) < 0.001,
    }

def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--N", type=int, default=80000)
    parser.add_argument("--p", type=int, required=True)
    parser.add_argument("--K", type=int, default=100)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()
    
    print(f"BMCTE v2 NUVMAP sparse: N={args.N}, p={args.p}")
    print(f"Threshold check: p/N = {args.p/args.N:.6f} (target 0.142857)")
    
    result = run_bmcte_sparse(args.N, args.p, args.K, args.seed)
    RECEIPT_PATH.write_text(json.dumps(result, indent=2))
    
    print(f"Saved {result['nuvmap_states_saved']} states to NUVMAP")
    print(f"λ_theory={result['lambda_theory']:.2e}")
    return 0

if __name__ == "__main__":
    import sys
    sys.exit(main())