#!/usr/bin/env python3
from __future__ import annotations

import argparse
import platform
import time
from pathlib import Path
from typing import Any

import numpy as np

from runtime import InferenceSession, first_output, load_json, max_abs_delta, read_wav, write_json


SCRIPT_ROOT = Path(__file__).resolve().parents[1]
REPO_LAYOUT = (SCRIPT_ROOT / "model_convert/pulsar2/check0/axmodels").is_dir()
DEFAULT_MODELS_DIR = (
    SCRIPT_ROOT / "model_convert/pulsar2/check0/axmodels"
    if REPO_LAYOUT
    else SCRIPT_ROOT / "models/650"
)
DEFAULT_AUDIO_DIR = (
    SCRIPT_ROOT / "openWakeWord/tests/data"
    if REPO_LAYOUT
    else SCRIPT_ROOT / "audio/openwakeword"
)
DEFAULT_REFERENCE = (
    SCRIPT_ROOT / "model_convert/local_inference/results.json"
    if REPO_LAYOUT
    else SCRIPT_ROOT / "reference/local_inference_results.json"
)
DEFAULT_MEL_WEIGHTS = (
    SCRIPT_ROOT / "model_convert/board_assets/openwakeword_mel_weights.npz"
    if REPO_LAYOUT
    else SCRIPT_ROOT / "config/openwakeword_mel_weights.npz"
)
DEFAULT_OUTPUT = (
    SCRIPT_ROOT / "model_convert/board_inference/openwakeword_board.json"
    if REPO_LAYOUT
    else SCRIPT_ROOT / "outputs/openwakeword_ax650_board.json"
)
CLASSIFIERS = (
    "alexa_v0.1",
    "hey_jarvis_v0.1",
    "hey_mycroft_v0.1",
    "hey_rhasspy_v0.1",
    "timer_v0.1",
    "weather_v0.1",
)
EXPECTED = {
    "alexa_test.wav": "alexa_v0.1",
    "hey_mycroft_test.wav": "hey_mycroft_v0.1",
}


def mel_host_postprocess(value: np.ndarray) -> np.ndarray:
    clipped = np.asarray(value, dtype=np.float32)
    decibels = np.log(clipped).astype(np.float32)
    decibels = decibels * np.float32(10.0)
    decibels = decibels / np.float32(2.3025851249694824)
    minimum = np.max(decibels).astype(np.float32) - np.float32(80.0)
    return np.maximum(decibels, minimum).astype(np.float32)


def load_mel_weights(path: Path) -> dict[str, np.ndarray]:
    if not path.is_file():
        raise FileNotFoundError(path)
    with np.load(path) as values:
        result = {
            "real": np.asarray(values["real"], dtype=np.float32),
            "imag": np.asarray(values["imag"], dtype=np.float32),
            "mel": np.asarray(values["mel"], dtype=np.float32),
            "floor": np.asarray(values["floor"], dtype=np.float32).reshape(()),
        }
    expected_shapes = {
        "real": (257, 512),
        "imag": (257, 512),
        "mel": (257, 32),
    }
    for name, shape in expected_shapes.items():
        if result[name].shape != shape:
            raise ValueError(f"unexpected mel weight shape for {name}: {result[name].shape}")
    return result


def numpy_melspectrogram(
    value: np.ndarray, weights: dict[str, np.ndarray]
) -> np.ndarray:
    samples = np.asarray(value, dtype=np.float32).reshape(-1)
    if samples.shape != (1760,):
        raise ValueError(f"expected 1760 mel input samples, got {samples.shape}")
    frames = np.stack(
        [samples[start : start + 512] for start in range(0, 1280, 160)],
        axis=0,
    )
    real = frames @ weights["real"].T
    imag = frames @ weights["imag"].T
    power = np.square(real) + np.square(imag)
    mel = power @ weights["mel"]
    mel = np.maximum(mel, weights["floor"]).astype(np.float32)
    return mel[None, None, :, :]


def model_path(models_dir: Path, name: str, backend: str) -> Path:
    if backend == "axengine":
        return models_dir / f"openwakeword__{name}.axmodel"
    return models_dir / f"{name}.onnx"


def pad_chunks(samples: np.ndarray, chunk_size: int = 1280) -> np.ndarray:
    remainder = samples.size % chunk_size
    if remainder == 0:
        return samples
    return np.pad(samples, (0, chunk_size - remainder))


def load_sessions(
    models_dir: Path, backend: str, mel_backend: str
) -> dict[str, InferenceSession]:
    names = ("embedding_model",) + CLASSIFIERS
    if mel_backend == "model":
        names = ("melspectrogram",) + names
    sessions = {
        name: InferenceSession(model_path(models_dir, name, backend), backend)
        for name in names
    }
    for name, session in sessions.items():
        print(
            f"loaded {name}: "
            f"inputs={[(value.name, value.shape, str(value.dtype)) for value in session.inputs]} "
            f"outputs={[(value.name, value.shape, str(value.dtype)) for value in session.outputs]}",
            flush=True,
        )
    return sessions


def reference_by_audio(path: Path | None) -> dict[str, dict[str, Any]]:
    if path is None or not path.is_file():
        return {}
    result = load_json(path)
    return {
        Path(row["audio"]).name: row
        for row in result["openwakeword"]["clips"]
    }


def infer_clip(
    audio_path: Path,
    sessions: dict[str, InferenceSession],
    reference: dict[str, Any] | None,
    threshold: float,
    score_tolerance: float,
    mel_backend: str,
    mel_weights: dict[str, np.ndarray] | None,
) -> dict[str, Any]:
    samples, sample_rate = read_wav(audio_path)
    if sample_rate != 16000:
        raise ValueError(f"expected 16 kHz audio: {audio_path} has {sample_rate}")
    samples = pad_chunks(samples)

    history = np.zeros(480, dtype=np.int16)
    mel_buffer = np.ones((76, 32), dtype=np.float32)
    feature_buffer = np.zeros((34, 96), dtype=np.float32)
    scores: dict[str, list[list[float]]] = {name: [] for name in CLASSIFIERS}

    for start in range(0, samples.size, 1280):
        chunk = samples[start : start + 1280]
        mel_input = np.concatenate((history, chunk)).astype(np.float32)[None, :]
        history = np.concatenate((history, chunk))[-480:].astype(np.int16, copy=False)

        if mel_backend == "numpy":
            if mel_weights is None:
                raise RuntimeError("NumPy mel backend requires mel weights")
            mel_output = numpy_melspectrogram(mel_input, mel_weights)
        else:
            mel_session = sessions["melspectrogram"]
            mel_output = first_output(
                mel_session,
                mel_session.run({mel_session.inputs[0].name: mel_input}),
            )
        mel_output = mel_host_postprocess(mel_output)
        spec = np.squeeze(mel_output).astype(np.float32) / 10.0 + 2.0
        if spec.ndim != 2 or spec.shape[1] != 32:
            raise RuntimeError(f"unexpected mel output shape: {spec.shape}")
        mel_buffer = np.vstack((mel_buffer, spec))[-970:]

        embedding_input = mel_buffer[-76:][None, :, :, None].astype(np.float32)
        embedding_session = sessions["embedding_model"]
        embedding_output = first_output(
            embedding_session,
            embedding_session.run(
                {embedding_session.inputs[0].name: embedding_input}
            ),
        )
        feature = np.asarray(embedding_output).reshape(-1, 96)[-1]
        feature_buffer = np.vstack((feature_buffer, feature))[-120:]

        for name in CLASSIFIERS:
            classifier = sessions[name]
            frame_count = classifier.inputs[0].shape[1]
            classifier_input = feature_buffer[-frame_count:][None, :, :].astype(
                np.float32
            )
            output = first_output(
                classifier,
                classifier.run({classifier.inputs[0].name: classifier_input}),
            )
            scores[name].append(np.asarray(output).reshape(-1).tolist())

    max_scores = {
        name: np.max(np.asarray(values, dtype=np.float64), axis=0).tolist()
        for name, values in scores.items()
    }
    expected_model = EXPECTED.get(audio_path.name)
    expected_score = (
        float(np.max(max_scores[expected_model])) if expected_model is not None else None
    )
    expected_detected = (
        expected_score >= threshold if expected_score is not None else None
    )

    reference_deltas = None
    maximum_reference_delta = None
    reference_match = None
    if reference is not None:
        reference_scores = {
            name.removesuffix(".onnx"): value
            for name, value in reference["static_max_scores"].items()
        }
        reference_deltas = {
            name: max_abs_delta(reference_scores[name], max_scores[name])
            for name in CLASSIFIERS
        }
        maximum_reference_delta = max(reference_deltas.values())
        reference_match = maximum_reference_delta <= score_tolerance

    functional_pass = expected_detected is not False
    passed = functional_pass and reference_match is not False
    return {
        "audio": str(audio_path),
        "sample_rate": sample_rate,
        "sample_count": int(samples.size),
        "frame_count": int(samples.size // 1280),
        "max_scores": max_scores,
        "expected_model": expected_model,
        "threshold": threshold,
        "expected_score": expected_score,
        "expected_detected": expected_detected,
        "reference_deltas": reference_deltas,
        "maximum_reference_delta": maximum_reference_delta,
        "score_tolerance": score_tolerance,
        "reference_match": reference_match,
        "passed": passed,
    }


def main(
    default_models_dir: Path = DEFAULT_MODELS_DIR,
    default_output: Path = DEFAULT_OUTPUT,
    target_hardware: str = "AX650",
) -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=("axengine", "onnx"), default="axengine")
    parser.add_argument("--mel-backend", choices=("numpy", "model"), default="numpy")
    parser.add_argument("--mel-weights", type=Path, default=DEFAULT_MEL_WEIGHTS)
    parser.add_argument("--models-dir", type=Path, default=default_models_dir)
    parser.add_argument("--audio-dir", type=Path, default=DEFAULT_AUDIO_DIR)
    parser.add_argument(
        "--reference",
        type=Path,
        default=DEFAULT_REFERENCE,
    )
    parser.add_argument(
        "--output", type=Path, default=default_output
    )
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--score-tolerance", type=float, default=0.15)
    args = parser.parse_args()

    audio_files = sorted(args.audio_dir.glob("*.wav"))
    if not audio_files:
        raise FileNotFoundError(f"no WAV files found under {args.audio_dir}")

    references = reference_by_audio(args.reference)
    mel_weights = load_mel_weights(args.mel_weights) if args.mel_backend == "numpy" else None
    sessions = load_sessions(args.models_dir, args.backend, args.mel_backend)
    started = time.perf_counter()
    clips = []
    for path in audio_files:
        clip = infer_clip(
            path,
            sessions,
            references.get(path.name),
            args.threshold,
            args.score_tolerance,
            args.mel_backend,
            mel_weights,
        )
        clips.append(clip)
        print(
            f"{path.name}: expected={clip['expected_model']} "
            f"score={clip['expected_score']} detected={clip['expected_detected']} "
            f"reference_delta={clip['maximum_reference_delta']} "
            f"passed={clip['passed']}",
            flush=True,
        )

    result = {
        "backend": args.backend,
        "target_hardware": target_hardware,
        "platform": platform.platform(),
        "models_dir": str(args.models_dir),
        "mel_backend": args.mel_backend,
        "mel_weights": str(args.mel_weights) if mel_weights is not None else None,
        "audio_dir": str(args.audio_dir),
        "reference": str(args.reference) if args.reference else None,
        "elapsed_seconds": time.perf_counter() - started,
        "clips": clips,
        "all_passed": all(clip["passed"] for clip in clips),
    }
    write_json(args.output, result)
    print(f"wrote {args.output}")
    print(f"all_passed={result['all_passed']}")
    if not result["all_passed"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()