import argparse
from pathlib import Path

import numpy as np
import onnxruntime as ort

from common import extract_mel_features, get_chunk_output_length_value, run_chunked_torch, split_mel_features


def parse_args():
    parser = argparse.ArgumentParser(description="Run Qwen3-ASR audio encoder ONNX by chunk and optionally compare to PyTorch.")
    parser.add_argument("--model-path", type=str, default=".", help="Path to Qwen3-ASR model directory.")
    parser.add_argument(
        "--onnx-path",
        type=str,
        default="rknn_deploy/audio_encoder/onnx/qwen3_asr_audio_chunk100.onnx",
        help="Path to exported ONNX.",
    )
    parser.add_argument("--audio-path", type=str, required=True, help="Path to input audio.")
    parser.add_argument("--chunk-frames", type=int, default=100, help="Fixed mel chunk length.")
    parser.add_argument("--compare-torch", action="store_true", help="Compare concatenated ONNX output with PyTorch.")
    parser.add_argument("--savepath", type=str, default=None, help="Optional path to save concatenated features as .npy.")
    return parser.parse_args()


def main():
    args = parse_args()
    input_features, feature_len = extract_mel_features(args.model_path, args.audio_path)
    chunks = split_mel_features(input_features, feature_len, args.chunk_frames)

    session = ort.InferenceSession(args.onnx_path, providers=["CPUExecutionProvider"])

    outputs = []
    for chunk, chunk_len in chunks:
        session_outputs = session.run(
            None,
            {
                "input_features": chunk[None, ...].astype(np.float32),
                "feature_len": np.asarray([chunk_len], dtype=np.int32),
            },
        )
        ort_features = session_outputs[0]
        if len(session_outputs) >= 2:
            valid_len = int(np.asarray(session_outputs[1]).reshape(-1)[0])
        else:
            valid_len = get_chunk_output_length_value(chunk_len)
        outputs.append(ort_features[0, :valid_len])

    concat_features = np.concatenate(outputs, axis=0) if outputs else np.zeros((0, 2048), dtype=np.float32)
    print(f"input_feature_len: {feature_len}")
    print(f"num_chunks: {len(chunks)}")
    print(f"audio_features: {concat_features.shape}")

    if args.savepath:
        savepath = Path(args.savepath)
        savepath.parent.mkdir(parents=True, exist_ok=True)
        np.save(savepath, concat_features)
        print(f"saved_features: {savepath}")

    if args.compare_torch:
        ref = run_chunked_torch(
            model_path=args.model_path,
            input_features=input_features,
            feature_len=feature_len,
            chunk_frames=args.chunk_frames,
        )
        max_diff = float(np.max(np.abs(ref - concat_features))) if ref.size and concat_features.size else 0.0
        print(f"torch_features: {ref.shape}")
        print(f"max_abs_diff(torch_vs_onnx): {max_diff:.8f}")


if __name__ == "__main__":
    main()