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

def research_metrology(iter_name="16_resolution_independent", model="awrl1432"):
    project_root = Path(__file__).parent.parent.parent
    base_path = project_root / "Shenron_debug" / "iterations" / iter_name / model
    files = sorted(list(base_path.glob("*.npy")))
    
    ranges = []
    max_mags = []
    all_mags = []

    print(f"🔬 Researching Physical Trends for {iter_name}...")

    # Sample frames across the simulation (1-250)
    for f in files[::10]:
        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]
        
        # Target Tracking (Assume the brightest point is the Ego-relative lead car)
        if len(rho) > 0:
            idx = np.argmax(mags)
            ranges.append(rho[idx])
            max_mags.append(mags[idx])
            all_mags.extend(mags.tolist())

    # 1. Range-Magnitude Analysis
    ranges = np.array(ranges)
    max_mags = np.array(max_mags)
    
    # Fit to 1/R^2 (since tx_loss was 2, and rx_loss is usually handled in radar eq)
    # Actually, let's see what the actual decay is.
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    plt.scatter(ranges, max_mags, alpha=0.6, label='Detected Peaks')
    plt.title("Physical Range-Magnitude Decay")
    plt.xlabel("Range (m)")
    plt.ylabel("Magnitude")
    plt.grid(True, alpha=0.3)
    
    # 2. SNR Distribution (Histogram)
    plt.subplot(1, 2, 2)
    plt.hist(all_mags, bins=50, color='skyblue', edgecolor='black', alpha=0.7)
    plt.yscale('log')
    plt.title("Magnitude Distribution (SNR Separation)")
    plt.xlabel("Magnitude Value")
    plt.ylabel("Point Count (Log)")
    plt.grid(True, alpha=0.3, axis='y')

    plt.tight_layout()
    plt.savefig("metrology_research.png")
    print("✅ Research plots saved to metrology_research.png")

if __name__ == "__main__":
    project_root = Path(__file__).parent.parent.parent
    os.chdir(str(project_root)) # Switch to root for data path consistency
    research_metrology()