"""Check scaled fixed-point FFT against a floating NumPy reference and estimate range."""
import argparse
import csv
import json
from pathlib import Path
import numpy as np


def load_reference(path):
    with np.load(path) as data:
        return data['fft'], json.loads(str(data['config']))


def read_simulation(path):
    frames = {}
    with Path(path).open(newline='') as f:
        for row in csv.reader(f):
            frame, *values = map(int, row)
            frames.setdefault(frame, []).append(values)
    if sorted(frames) != [1, 2]:
        raise ValueError('Expected simulated frames 1 and 2')
    return frames


def check_frame(rows, reference, config, case_id):
    nfft = config['fft_length']
    if len(rows) != nfft or [r[0] for r in rows] != list(range(nfft)):
        raise ValueError('Expected exactly 512 bins in natural order')
    values = np.array(rows, dtype=np.int64)
    i, q, power = values[:, 1], values[:, 2], values[:, 3]
    if np.any(i < -32768) or np.any(i > 32767) or np.any(q < -32768) or np.any(q > 32767):
        raise ValueError('Invalid signed 16-bit FFT output')
    if not np.array_equal(power, i*i + q*q):
        raise ValueError('Power pipeline disagrees with complex output')
    error = float(max(np.max(np.abs(i-reference.real)), np.max(np.abs(q-reference.imag))))
    # Stage rounding makes NumPy approximate; do not describe this as bit exact.
    if error > 8:
        raise ValueError(f'FFT error {error:.3f} exceeds 8 Q15 integer-count tolerance')
    peak = int(np.argmax(power[1:nfft//2])) + 1
    expected = int(np.argmax(np.abs(reference[1:nfft//2])**2)) + 1
    if peak != expected:
        raise ValueError(f'Expected peak {expected}, got {peak}')
    fs = config['sample_rate_hz']
    slope = config['bandwidth_hz'] / config['chirp_duration_s']
    spacing = config['c_m_s'] * fs / (2*slope*nfft)
    truth = config['cases'][case_id]['range_m']
    estimated = peak * spacing
    if abs(estimated - truth) > spacing / 2:
        raise ValueError('Range error exceeds half a bin for this noiseless single target')
    return {'case_id': case_id, 'bins': nfft, 'peak_bin': peak,
            'peak_power_q30': int(power[peak]), 'peak_power': float(power[peak]/2**30),
            'max_component_error_q15_counts': error, 'beat_frequency_hz': peak*fs/nfft,
            'range_bin_spacing_m': spacing, 'true_range_m': truth,
            'estimated_range_m': estimated, 'range_error_m': estimated-truth}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--csv', required=True, type=Path)
    parser.add_argument('--reference', required=True, type=Path)
    parser.add_argument('--out', required=True, type=Path)
    args = parser.parse_args()
    reference, config = load_reference(args.reference)
    checked = {str(f): check_frame(rows, reference[f-1], config, f-1)
               for f, rows in read_simulation(args.csv).items()}
    args.out.write_text(json.dumps({'verified': True, 'frames': checked}, indent=2), encoding='utf-8')
    print('FFT_NUMPY_PASS', json.dumps(checked))


if __name__ == '__main__':
    main()