#!/usr/bin/env python3
"""Export Paraformer to GGUF with experimental Q5_0/Q4_0 support."""

import argparse
import glob
import json
import os
import re

import gguf
import numpy as np
import torch


def parse_mvn(path):
    blocks = [
        np.array([float(x) for x in block.split()], np.float32)
        for block in re.findall(r"\[([^\]]*)\]", open(path).read())
    ]
    vectors = [block for block in blocks if block.size > 1]
    return vectors[0], vectors[1]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pt", required=True)
    parser.add_argument("--mvn", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument(
        "--wtype",
        default="q5_0",
        choices=["f32", "f16", "q8_0", "q5_0", "q4_0"],
    )
    parser.add_argument("--tokens", default=None)
    args = parser.parse_args()

    state_dict = torch.load(args.model_pt, map_location="cpu")
    state_dict = state_dict.get("state_dict", state_dict)
    writer = gguf.GGUFWriter(args.out, "paraformer")
    writer.add_uint32("pf.enc.output_size", 512)
    writer.add_uint32("pf.enc.attention_heads", 4)
    writer.add_uint32("pf.enc.num_blocks", 50)
    writer.add_uint32("pf.enc.kernel_size", 11)
    writer.add_uint32("pf.dec.num_blocks", 16)
    writer.add_uint32("pf.dec.att_layer_num", 16)
    writer.add_uint32("pf.dec.decoders3", 1)
    writer.add_uint32("pf.dec.attention_heads", 4)
    writer.add_uint32("pf.dec.kernel_size", 11)
    writer.add_uint32("pf.vocab_size", 8404)

    token_path = args.tokens or (
        glob.glob(os.path.join(os.path.dirname(args.model_pt), "tokens.json")) + [None]
    )[0]
    if token_path and os.path.exists(token_path):
        with open(token_path, encoding="utf-8") as token_file:
            tokens = json.load(token_file)
        writer.add_array("pf.vocab", tokens)
        print(f"embedded pf.vocab ({len(tokens)} tokens) from {token_path}")
    else:
        print("WARNING: tokens.json not found - GGUF will have no vocabulary")

    writer.add_float32("pf.predictor.tail_threshold", 0.45)
    writer.add_float32("pf.predictor.threshold", 1.0)
    shift, scale = parse_mvn(args.mvn)
    writer.add_tensor("cmvn.shift", shift)
    writer.add_tensor("cmvn.scale", scale)

    from gguf import GGMLQuantizationType as QuantType
    from gguf import quants

    quant_types = {
        "q8_0": QuantType.Q8_0,
        "q5_0": QuantType.Q5_0,
        "q4_0": QuantType.Q4_0,
    }
    quant_type = quant_types.get(args.wtype)
    quant_block_size = (
        quants._type_traits[quant_type].block_size if quant_type is not None else 1
    )
    tensor_count = 0
    quantized_count = 0

    for name, value in state_dict.items():
        if not name.startswith(("encoder.", "decoder.", "predictor.")):
            continue
        if name == "decoder.embed.0.weight":
            continue

        array = value.detach().to(torch.float32).contiguous().numpy()
        if name.endswith("fsmn_block.weight") and array.ndim == 3:
            array = np.ascontiguousarray(array[:, 0, :].T)
        elif (
            args.wtype == "f16"
            and array.ndim == 2
            and "norm" not in name
            and "cif_output" not in name
        ):
            array = array.astype(np.float16)

        can_quantize = (
            quant_type is not None
            and array.ndim == 2
            and "norm" not in name
            and "fsmn_block" not in name
            and "predictor" not in name
            and array.shape[1] % quant_block_size == 0
        )
        if can_quantize:
            writer.add_tensor(
                name,
                quants.quantize(array, quant_type),
                raw_dtype=quant_type,
            )
            quantized_count += 1
        else:
            writer.add_tensor(name, array)
        tensor_count += 1

    print(
        f"writing {tensor_count} tensors (+cmvn), {quantized_count} quantized "
        f"as {args.wtype}, to {args.out}"
    )
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()
    print(f"done: {args.out} ({os.path.getsize(args.out) / 1e6:.1f} MB)")


if __name__ == "__main__":
    main()