#!/usr/bin/env python3
"""Compare teacher, FP32 student, and INT8 student embeddings on real text."""

from __future__ import annotations

import argparse
import json
import math
import sys
import time
from pathlib import Path

import numpy as np
import onnxruntime as ort

from vosk_booktts_experiment import WordPieceTokenizer, load_corpus, normalize_text


def cosine_rows(left: np.ndarray, right: np.ndarray) -> np.ndarray:
    numerator = np.sum(left * right, axis=-1)
    denominator = np.linalg.norm(left, axis=-1) * np.linalg.norm(right, axis=-1)
    return numerator / np.maximum(denominator, 1e-12)


def compare(args: argparse.Namespace) -> None:
    tokenizer = WordPieceTokenizer(args.vocab)
    token_map = np.load(args.token_map)
    compact_unk = int(token_map[tokenizer.vocabulary["[UNK]"]])
    sessions = {
        "teacher": ort.InferenceSession(str(args.teacher), providers=["CPUExecutionProvider"]),
        "student_fp32": ort.InferenceSession(str(args.student_fp32), providers=["CPUExecutionProvider"]),
        "student_int8": ort.InferenceSession(str(args.student_int8), providers=["CPUExecutionProvider"]),
    }
    timings = {name: 0.0 for name in sessions}
    teacher_cosines: list[np.ndarray] = []
    quantized_cosines: list[np.ndarray] = []
    teacher_squared_error = 0.0
    quantized_squared_error = 0.0
    elements = 0
    tokens = 0
    unseen_tokens = 0
    phrases = load_corpus(args.corpus)[: args.limit or None]
    for text in phrases:
        ids, _ = tokenizer.encode(normalize_text(text))
        shape = (1, ids.size)
        feeds = {
            "input_ids": ids.reshape(shape),
            "attention_mask": np.ones(shape, dtype=np.int64),
            "token_type_ids": np.zeros(shape, dtype=np.int64),
        }
        outputs: dict[str, np.ndarray] = {}
        for name, session in sessions.items():
            started = time.perf_counter()
            outputs[name] = session.run(None, feeds)[0].astype(np.float32, copy=False)
            timings[name] += time.perf_counter() - started
        teacher = outputs["teacher"]
        fp32 = outputs["student_fp32"]
        int8 = outputs["student_int8"]
        if teacher.shape != fp32.shape or fp32.shape != int8.shape:
            raise ValueError(f"Output shape mismatch for {text!r}: {teacher.shape}, {fp32.shape}, {int8.shape}")
        teacher_cosines.append(cosine_rows(teacher, fp32))
        quantized_cosines.append(cosine_rows(fp32, int8))
        teacher_squared_error += float(np.sum((teacher - fp32) ** 2))
        quantized_squared_error += float(np.sum((fp32 - int8) ** 2))
        elements += teacher.size
        tokens += ids.size
        unseen_tokens += int(np.sum((token_map[ids] == compact_unk) & (ids != tokenizer.vocabulary["[UNK]"])))
    teacher_cosine = np.concatenate(teacher_cosines)
    quantized_cosine = np.concatenate(quantized_cosines)
    report = {
        "phrases": len(phrases),
        "tokens": tokens,
        "unseen_tokens": unseen_tokens,
        "unseen_percent": 100.0 * unseen_tokens / tokens,
        "teacher_vs_student_fp32": {
            "mean_token_cosine": float(np.mean(teacher_cosine)),
            "p05_token_cosine": float(np.percentile(teacher_cosine, 5)),
            "mse": teacher_squared_error / elements,
        },
        "student_fp32_vs_int8": {
            "mean_token_cosine": float(np.mean(quantized_cosine)),
            "p05_token_cosine": float(np.percentile(quantized_cosine, 5)),
            "mse": quantized_squared_error / elements,
        },
        "total_inference_seconds": timings,
        "milliseconds_per_phrase": {name: seconds * 1000.0 / len(phrases) for name, seconds in timings.items()},
    }
    if not all(math.isfinite(value) for value in (teacher_squared_error, quantized_squared_error)):
        raise ValueError("Non-finite comparison result")
    rendered = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
    if args.output:
        if args.output.exists():
            raise FileExistsError(f"Refusing to overwrite {args.output}")
        args.output.write_text(rendered, encoding="utf-8")
    print(rendered, end="")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--teacher", type=Path, required=True)
    parser.add_argument("--student-fp32", type=Path, required=True)
    parser.add_argument("--student-int8", type=Path, required=True)
    parser.add_argument("--vocab", type=Path, required=True)
    parser.add_argument("--token-map", type=Path, required=True)
    parser.add_argument("--corpus", type=Path, required=True)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--limit", type=int, default=0)
    return parser.parse_args()


if __name__ == "__main__":
    sys.exit(compare(parse_args()) or 0)