#!/usr/bin/env python3
"""
Test script to run and compare both PyTorch and LiteRT / TFLite inference.
Synthesizes a realistic audio chirp sample if no test WAV is provided.
"""

import argparse
import json
from pathlib import Path
import numpy as np
import scipy.io.wavfile as wavfile
import tensorflow as tf

from src.inference import predict_wav, identify, _summed, genus_of, group_of, DEFAULT_THRESHOLD
from src.audio import read_audio, audio_quality, resample, crop_offsets, logmel, SR

ROOT = Path(__file__).resolve().parent


def generate_synthetic_chirp(path, duration=3.0, sr=16000):
    """Generate a 3-second synthetic insect chirp WAV for local test inference."""
    t = np.linspace(0, duration, int(sr * duration), endpoint=False)
    # Carrier frequency around 4500 Hz modulated with bursts (simulating cricket/cicada pulse train)
    carrier = np.sin(2 * np.pi * 4500 * t)
    pulse = (np.sin(2 * np.pi * 25 * t) > 0.3).astype(np.float32)
    waveform = (carrier * pulse * 0.4).astype(np.float32)
    # Save as 16-bit PCM WAV
    wavfile.write(str(path), sr, (waveform * 32767).astype(np.int16))
    return path


def predict_tflite(path, model_dir=ROOT, tflite_path=None):
    """Run inference using the LiteRT / TFLite model."""
    model_dir = Path(model_dir)
    cfg = json.loads((model_dir / 'config.json').read_text())
    labels = json.loads((model_dir / 'labels.json').read_text())

    if tflite_path is None:
        tflite_path = model_dir / 'model.tflite'
    else:
        tflite_path = Path(tflite_path)

    if not tflite_path.is_file():
        raise FileNotFoundError(f"LiteRT model not found at '{tflite_path}'. Run export_litert.py first.")

    interpreter = tf.lite.Interpreter(model_path=str(tflite_path))
    interpreter.allocate_tensors()
    in_idx = interpreter.get_input_details()[0]['index']
    out_idx = interpreter.get_output_details()[0]['index']

    sr, y = read_audio(path)
    ok, reasons = audio_quality(y)
    y16 = resample(y, sr)
    offs = crop_offsets(len(y16), int(cfg.get('eval_crops', 1)))

    logits_list = []
    for off in offs:
        feat = logmel(y16, SR, off)  # (1, 64, 201)
        feat_batch = feat[None, :, :, :].astype(np.float32)  # (1, 1, 64, 201)
        interpreter.set_tensor(in_idx, feat_batch)
        interpreter.invoke()
        logits = interpreter.get_tensor(out_idx)  # (1, 32)
        logits_list.append(logits[0])

    # Softmax over logits, then mean across crops
    logits_arr = np.array(logits_list)
    exp_logits = np.exp(logits_arr - np.max(logits_arr, axis=-1, keepdims=True))
    crop_probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
    probs = crop_probs.mean(axis=0)

    order = np.argsort(probs)[::-1][:5]
    thr = float(cfg.get('confidence_threshold', DEFAULT_THRESHOLD))

    return {
        'identification': identify(probs, labels, thr),
        'genus_scores': [{'genus': g, 'score': v} for g, v in _summed(probs, labels, genus_of)[:3]],
        'group_scores': [{'group': g, 'score': v} for g, v in _summed(probs, labels, group_of)],
        'predicted_species': labels[str(int(order[0]))],
        'confidence': float(probs[order[0]]),
        'top_k': [{'species': labels[str(int(i))], 'score': float(probs[i])} for i in order],
        'quality_pass': bool(ok),
        'quality_reasons': reasons,
        'model_version': cfg['model_version'],
        'runtime': 'LiteRT / TFLite'
    }


def main():
    parser = argparse.ArgumentParser(description="Test PyTorch and LiteRT insect classification inference")
    parser.add_argument('--wav', default=None, help="Path to input audio file (.wav). If not provided, a sample chirp is created.")
    parser.add_argument('--model-dir', default='.', help="Model directory")
    args = parser.parse_args()

    model_dir = Path(args.model_dir)
    wav_path = args.wav
    if wav_path is None:
        wav_path = model_dir / 'test_sample.wav'
        print(f"[*] Generating synthetic 3-second insect test chirp at '{wav_path}'...")
        generate_synthetic_chirp(wav_path)
    else:
        wav_path = Path(wav_path)

    print(f"\n[1] Running PyTorch Inference on '{wav_path}'...")
    pt_res = predict_wav(wav_path, model_dir)
    print(f"    PyTorch Top-1 Species : {pt_res['predicted_species']}")
    print(f"    PyTorch Confidence    : {pt_res['confidence']:.4f}")
    print(f"    PyTorch Identification: {pt_res['identification']}")

    print(f"\n[2] Running LiteRT / TFLite Inference on '{wav_path}'...")
    tflite_res = predict_tflite(wav_path, model_dir)
    print(f"    LiteRT Top-1 Species  : {tflite_res['predicted_species']}")
    print(f"    LiteRT Confidence     : {tflite_res['confidence']:.4f}")
    print(f"    LiteRT Identification : {tflite_res['identification']}")

    # Parity check
    species_match = pt_res['predicted_species'] == tflite_res['predicted_species']
    conf_diff = abs(pt_res['confidence'] - tflite_res['confidence'])
    print("\n" + "=" * 60)
    print(" VERIFICATION SUMMARY:")
    print(f" - Top Predicted Species Match : {'YES' if species_match else 'NO'} ({pt_res['predicted_species']})")
    print(f" - Confidence Difference       : {conf_diff:.8e}")
    print(f" - Audio Quality Check Passed  : {pt_res['quality_pass']}")
    print("=" * 60)

    if species_match and conf_diff < 1e-4:
        print("[+] Test PASSED: Both runtimes agree with exact parity.\n")
    else:
        print("[-] Test WARNING: Slight divergence detected.\n")


if __name__ == '__main__':
    main()