import argparse
import time
import numpy as np
import onnx
from onnxsim import simplify
import onnxruntime as ort
import onnxoptimizer
import torch
from model_onnx import SynthesizerTrn
import utils
from hubert import hubert_model_onnx

def main(HubertExport,NetExport):

    path = "NyaruTaffy"

    if(HubertExport):
        device = torch.device("cuda")
        hubert_soft = utils.get_hubert_model()
        test_input = torch.rand(1, 1, 16000)
        input_names = ["source"]
        output_names = ["embed"]
        torch.onnx.export(hubert_soft.to(device),
                        test_input.to(device),
                        "hubert3.0.onnx",
                        dynamic_axes={
                            "source": {
                                2: "sample_length"
                            }
                        },
                        verbose=False,
                        opset_version=13,
                        input_names=input_names,
                        output_names=output_names)
    if(NetExport):
        device = torch.device("cuda")
        hps = utils.get_hparams_from_file(f"checkpoints/{path}/config.json")
        SVCVITS = SynthesizerTrn(
            hps.data.filter_length // 2 + 1,
            hps.train.segment_size // hps.data.hop_length,
            **hps.model)
        _ = utils.load_checkpoint(f"checkpoints/{path}/model.pth", SVCVITS, None)
        _ = SVCVITS.eval().to(device)
        for i in SVCVITS.parameters():
            i.requires_grad = False
        test_hidden_unit = torch.rand(1, 50, 256)
        test_lengths = torch.LongTensor([50])
        test_pitch = torch.rand(1, 50)
        test_sid = torch.LongTensor([0])
        input_names = ["hidden_unit", "lengths", "pitch", "sid"]
        output_names = ["audio", ]
        SVCVITS.eval()
        torch.onnx.export(SVCVITS,
                        (
                            test_hidden_unit.to(device),
                            test_lengths.to(device),
                            test_pitch.to(device),
                            test_sid.to(device)
                        ),
                        f"checkpoints/{path}/model.onnx",
                        dynamic_axes={
                            "hidden_unit": [0, 1],
                            "pitch": [1]
                        },
                        do_constant_folding=False,
                        opset_version=16,
                        verbose=False,
                        input_names=input_names,
                        output_names=output_names)


if __name__ == '__main__':
    main(False,True)