#!/usr/bin/env python3
"""
Co-simulation Comparison: RTL vs Python Model for AERIS-10 Matched Filter.

Compares the RTL matched filter output (from tb_mf_cosim.v) against the
Python model golden reference (from gen_mf_cosim_golden.py).

Two modes of operation:
  1. Synthesis branch (no -DSIMULATION): RTL uses fft_engine.v with fixed-point
     twiddle ROM (fft_twiddle_1024.mem) and frequency_matched_filter.v. The
     Python model was built to match this exactly. Expect BIT-PERFECT results
     (correlation = 1.0, energy ratio = 1.0).

  2. SIMULATION branch (-DSIMULATION): RTL uses behavioral FFT with floating-
     point twiddles ($rtoi($cos*32767)) and shift-then-add conjugate multiply.
     Python model uses fixed-point twiddles and add-then-round. Expect large
     numerical differences; only state-machine mechanics are validated.

Usage:
    python3 compare_mf.py [scenario|all]

    scenario: chirp, dc, impulse, tone5 (default: chirp)
    all: run all scenarios

Author: Phase 0.5 matched-filter co-simulation suite for PLFM_RADAR
"""

import math
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


# =============================================================================
# Configuration
# =============================================================================

FFT_SIZE = 1024

SCENARIOS = {
    'chirp': {
        'golden_csv': 'mf_golden_py_chirp.csv',
        'rtl_csv': 'rtl_mf_chirp.csv',
        'description': 'Radar chirp: 2 targets vs ref chirp',
    },
    'dc': {
        'golden_csv': 'mf_golden_py_dc.csv',
        'rtl_csv': 'rtl_mf_dc.csv',
        'description': 'DC autocorrelation (I=0x1000)',
    },
    'impulse': {
        'golden_csv': 'mf_golden_py_impulse.csv',
        'rtl_csv': 'rtl_mf_impulse.csv',
        'description': 'Impulse autocorrelation (delta at n=0)',
    },
    'tone5': {
        'golden_csv': 'mf_golden_py_tone5.csv',
        'rtl_csv': 'rtl_mf_tone5.csv',
        'description': 'Tone autocorrelation (bin 5, amp=8000)',
    },
}

# Thresholds for pass/fail
# These are generous because of the fundamental twiddle arithmetic differences
# between the SIMULATION branch (float twiddles) and Python model (fixed twiddles)
ENERGY_CORR_MIN = 0.80       # Min correlation of magnitude spectra
TOP_PEAK_OVERLAP_MIN = 0.50  # At least 50% of top-N peaks must overlap
RMS_RATIO_MAX = 50.0         # Max ratio of RMS energies (generous, since gain differs)
ENERGY_RATIO_MIN = 0.001     # Min ratio (total energy RTL / total energy Python)
ENERGY_RATIO_MAX = 1000.0    # Max ratio


# =============================================================================
# Helper functions
# =============================================================================

def load_csv(filepath):
    """Load CSV with columns (bin, out_i/range_profile_i, out_q/range_profile_q)."""
    vals_i = []
    vals_q = []
    with open(filepath) as f:
        f.readline()  # Skip header
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split(',')
            vals_i.append(int(parts[1]))
            vals_q.append(int(parts[2]))
    return vals_i, vals_q


def magnitude_spectrum(vals_i, vals_q):
    """Compute magnitude = |I| + |Q| for each bin (L1 norm, matches RTL)."""
    return [abs(i) + abs(q) for i, q in zip(vals_i, vals_q, strict=False)]


def magnitude_l2(vals_i, vals_q):
    """Compute magnitude = sqrt(I^2 + Q^2) for each bin."""
    return [math.sqrt(i*i + q*q) for i, q in zip(vals_i, vals_q, strict=False)]


def total_energy(vals_i, vals_q):
    """Compute total energy (sum of I^2 + Q^2)."""
    return sum(i*i + q*q for i, q in zip(vals_i, vals_q, strict=False))


def rms_magnitude(vals_i, vals_q):
    """Compute RMS of complex magnitude."""
    n = len(vals_i)
    if n == 0:
        return 0.0
    return math.sqrt(sum(i*i + q*q for i, q in zip(vals_i, vals_q, strict=False)) / n)


def pearson_correlation(a, b):
    """Compute Pearson correlation coefficient between two lists."""
    n = len(a)
    if n < 2:
        return 0.0
    mean_a = sum(a) / n
    mean_b = sum(b) / n
    cov = sum((a[i] - mean_a) * (b[i] - mean_b) for i in range(n))
    std_a_sq = sum((x - mean_a) ** 2 for x in a)
    std_b_sq = sum((x - mean_b) ** 2 for x in b)
    if std_a_sq < 1e-10 or std_b_sq < 1e-10:
        return 1.0 if abs(mean_a - mean_b) < 1.0 else 0.0
    return cov / math.sqrt(std_a_sq * std_b_sq)


def find_peak(vals_i, vals_q):
    """Find the bin with the maximum L1 magnitude."""
    mags = magnitude_spectrum(vals_i, vals_q)
    peak_bin = 0
    peak_mag = mags[0]
    for i in range(1, len(mags)):
        if mags[i] > peak_mag:
            peak_mag = mags[i]
            peak_bin = i
    return peak_bin, peak_mag


def top_n_peaks(mags, n=10):
    """Find the top-N peak bins by magnitude. Returns set of bin indices."""
    indexed = sorted(enumerate(mags), key=lambda x: -x[1])
    return {idx for idx, _ in indexed[:n]}


def spectral_peak_overlap(mags_a, mags_b, n=10):
    """Fraction of top-N peaks from A that also appear in top-N of B."""
    peaks_a = top_n_peaks(mags_a, n)
    peaks_b = top_n_peaks(mags_b, n)
    if len(peaks_a) == 0:
        return 1.0
    overlap = peaks_a & peaks_b
    return len(overlap) / len(peaks_a)


# =============================================================================
# Comparison for one scenario
# =============================================================================

def compare_scenario(scenario_name, config, base_dir):
    """Compare one scenario. Returns (pass/fail, result_dict)."""

    golden_path = os.path.join(base_dir, config['golden_csv'])
    rtl_path = os.path.join(base_dir, config['rtl_csv'])

    if not os.path.exists(golden_path):
        return False, {}
    if not os.path.exists(rtl_path):
        return False, {}

    py_i, py_q = load_csv(golden_path)
    rtl_i, rtl_q = load_csv(rtl_path)


    if len(py_i) != FFT_SIZE or len(rtl_i) != FFT_SIZE:
        return False, {}

    # ---- Metric 1: Energy ----
    py_energy = total_energy(py_i, py_q)
    rtl_energy = total_energy(rtl_i, rtl_q)
    py_rms = rms_magnitude(py_i, py_q)
    rtl_rms = rms_magnitude(rtl_i, rtl_q)

    if py_energy > 0 and rtl_energy > 0:
        energy_ratio = rtl_energy / py_energy
        rms_ratio = rtl_rms / py_rms
    elif py_energy == 0 and rtl_energy == 0:
        energy_ratio = 1.0
        rms_ratio = 1.0
    else:
        energy_ratio = float('inf') if py_energy == 0 else 0.0
        rms_ratio = float('inf') if py_rms == 0 else 0.0


    # ---- Metric 2: Peak location ----
    py_peak_bin, _py_peak_mag = find_peak(py_i, py_q)
    rtl_peak_bin, _rtl_peak_mag = find_peak(rtl_i, rtl_q)


    # ---- Metric 3: Magnitude spectrum correlation ----
    py_mag = magnitude_l2(py_i, py_q)
    rtl_mag = magnitude_l2(rtl_i, rtl_q)
    mag_corr = pearson_correlation(py_mag, rtl_mag)


    # ---- Metric 4: Top-N peak overlap ----
    # Use L1 magnitudes for peak finding (matches RTL)
    py_mag_l1 = magnitude_spectrum(py_i, py_q)
    rtl_mag_l1 = magnitude_spectrum(rtl_i, rtl_q)
    peak_overlap_10 = spectral_peak_overlap(py_mag_l1, rtl_mag_l1, n=10)
    peak_overlap_20 = spectral_peak_overlap(py_mag_l1, rtl_mag_l1, n=20)


    # ---- Metric 5: I and Q channel correlation ----
    corr_i = pearson_correlation(py_i, rtl_i)
    corr_q = pearson_correlation(py_q, rtl_q)


    # ---- Pass/Fail Decision ----
    # The SIMULATION branch uses floating-point twiddles ($cos/$sin) while
    # the Python model uses the fixed-point twiddle ROM (matching synthesis).
    # These are fundamentally different FFT implementations. We do NOT expect
    # structural similarity (correlation, peak overlap) between them.
    #
    # What we CAN verify:
    # 1. Both produce non-trivial output (state machine completes)
    # 2. Output count is correct (1024 samples)
    # 3. Energy is in a reasonable range (not wildly wrong)
    #
    # The true bit-accuracy comparison will happen when the synthesis branch
    # is simulated (xsim on remote server) using the same fft_engine.v that
    # the Python model was built to match.

    checks = []

    # Check 1: Both produce output
    both_have_output = py_energy > 0 and rtl_energy > 0
    checks.append(('Both produce output', both_have_output))

    # Check 2: RTL produced expected sample count
    correct_count = len(rtl_i) == FFT_SIZE
    checks.append(('Correct output count (1024)', correct_count))

    # Check 3: Energy ratio within generous bounds
    # Allow very wide range since twiddle differences cause large gain variation
    energy_ok = ENERGY_RATIO_MIN < energy_ratio < ENERGY_RATIO_MAX
    checks.append((f'Energy ratio in bounds ({ENERGY_RATIO_MIN}-{ENERGY_RATIO_MAX})',
                    energy_ok))

    # Print checks
    all_pass = True
    for _name, passed in checks:
        if not passed:
            all_pass = False

    result = {
        'scenario': scenario_name,
        'py_energy': py_energy,
        'rtl_energy': rtl_energy,
        'energy_ratio': energy_ratio,
        'rms_ratio': rms_ratio,
        'py_peak_bin': py_peak_bin,
        'rtl_peak_bin': rtl_peak_bin,
        'mag_corr': mag_corr,
        'peak_overlap_10': peak_overlap_10,
        'peak_overlap_20': peak_overlap_20,
        'corr_i': corr_i,
        'corr_q': corr_q,
        'passed': all_pass,
    }

    # Write detailed comparison CSV
    compare_csv = os.path.join(base_dir, f'compare_mf_{scenario_name}.csv')
    with open(compare_csv, 'w') as f:
        f.write('bin,py_i,py_q,rtl_i,rtl_q,py_mag,rtl_mag,diff_i,diff_q\n')
        for k in range(FFT_SIZE):
            f.write(f'{k},{py_i[k]},{py_q[k]},{rtl_i[k]},{rtl_q[k]},'
                    f'{py_mag_l1[k]},{rtl_mag_l1[k]},'
                    f'{rtl_i[k]-py_i[k]},{rtl_q[k]-py_q[k]}\n')

    return all_pass, result


# =============================================================================
# Main
# =============================================================================

def main():
    base_dir = os.path.dirname(os.path.abspath(__file__))

    arg = sys.argv[1].lower() if len(sys.argv) > 1 else 'chirp'

    if arg == 'all':
        run_scenarios = list(SCENARIOS.keys())
    elif arg in SCENARIOS:
        run_scenarios = [arg]
    else:
        sys.exit(1)


    results = []
    for name in run_scenarios:
        passed, result = compare_scenario(name, SCENARIOS[name], base_dir)
        results.append((name, passed, result))

    # Summary


    all_pass = True
    for _name, passed, result in results:
        if not result:
            all_pass = False
        else:
            if not passed:
                all_pass = False

    if all_pass:
        pass
    else:
        pass

    sys.exit(0 if all_pass else 1)


if __name__ == '__main__':
    main()
