import numpy as np
import os
from pathlib import Path
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import pandas as pd

def power_law(r, a, n):
    return a * np.power(r, -n)

def analyze_iteration(iter_name, base_path, model="awrl1432"):
    path = base_path / "iterations" / iter_name / model
    files = sorted(list(path.glob("*.npy")))
    
    ranges = []
    peaks = []
    all_mags = []
    
    for f in files:
        data = np.load(f)
        if data.shape[0] == 0: continue
        
        # Calculate radial range
        rho = np.linalg.norm(data[:, 0:3], axis=1)
        mags = data[:, 4]
        
        # Lead Car Peak Estimation
        # We look for the brightest cluster beyond 10m to avoid ego-ghosts
        mask = (rho > 14) & (rho < 100)
        if np.any(mask):
            idx = np.argmax(mags[mask])
            ranges.append(rho[mask][idx])
            peaks.append(mags[mask][idx])
        
        all_mags.extend(mags.tolist())
        
    return np.array(ranges), np.array(peaks), np.array(all_mags)

def run_metrology_suite():
    project_root = Path(__file__).parent.parent.parent
    base = project_root / "Shenron_debug"
    artifacts = Path(r"C:\Users\rakadu1\.gemini\antigravity\brain\67913a3c-cbc2-4fba-87e3-88fbea20f043\artifacts")

    # 1. Extract Data for 14b vs 16
    r_14, p_14, m_14 = analyze_iteration("14b_normalization", base)
    r_16, p_16, m_16 = analyze_iteration("16_resolution_independent", base)

    print("📊 Deep Metrology: Statistical Evidence Gathering...")

    # 2. Plot Magnitude Distribution (SNR Separation)
    plt.figure(figsize=(15, 6))
    
    plt.subplot(1, 2, 1)
    plt.hist(m_14, bins=100, color='red', alpha=0.4, label='Iter 14b (Bandaged)', density=True)
    plt.hist(m_16, bins=100, color='green', alpha=0.4, label='Iter 16 (Physical)', density=True)
    plt.yscale('log')
    plt.title("SNR Peak Separation Comparison")
    plt.xlabel("Magnitude")
    plt.ylabel("Probability Density (Log)")
    plt.legend()
    plt.grid(True, alpha=0.2)

    # 3. Plot Range Decay (The R^4 Law)
    plt.subplot(1, 2, 2)
    plt.scatter(r_14, p_14, color='red', s=10, alpha=0.3, label='Iter 14b')
    plt.scatter(r_16, p_16, color='green', s=10, alpha=0.5, label='Iter 16')
    
    # Fit Iter 16 to Power Law
    try:
        popt, _ = curve_fit(power_law, r_16, p_16, p0=[1e6, 2.0])
        r_fit = np.linspace(min(r_16), max(r_16), 100)
        plt.plot(r_fit, power_law(r_fit, *popt), 'k--', linewidth=2, label=f'16 Fit: 1/R^{popt[1]:.2f}')
    except Exception as e:
        print(f"Fit failed: {e}")

    plt.title("Physical Range-Magnitude Decay")
    plt.xlabel("Range (m)")
    plt.ylabel("Magnitude")
    plt.yscale('log')
    plt.xscale('log')
    plt.legend()
    plt.grid(True, alpha=0.2)

    plt.tight_layout()
    plt.savefig(artifacts / "deep_metrology_comparison.png")
    
    # 4. Context Independence Proof
    # Focus on Frames 50-80 (Buildings) vs 100-130 (Open Road)
    # Actually, we'll just check the consistency across all ranges.
    
    # Numerical Comparison
    print("\n🔍 PHYSICAL POINTERS:")
    print(f"{'Metric':<25} | {'Iter 14b':<15} | {'Iter 16':<15}")
    print("-" * 60)
    print(f"{'Mean Magnitude':<25} | {np.mean(m_14):<15.2f} | {np.mean(m_16):<15.2f}")
    if len(p_16) > 0 and len(p_14) > 0:
        print(f"{'Peak Var (Consistency)':<25} | {np.std(p_14)/np.mean(p_14):<15.2f} | {np.std(p_16)/np.mean(p_16):<15.2f}")
    
    print("\n✅ Metrology artifacts generated.")

if __name__ == "__main__":
    run_metrology_suite()