from rknn.api import RKNN
import onnxruntime as ort
import os

platforms=["rk3588", "rk3576"]

def search_for_models_decoder(base_dir):
    results = []

    for d1 in os.listdir(base_dir):
        if d1.startswith('.'):
            continue
        p1 = os.path.join(base_dir, d1)
        if not os.path.isdir(p1):
            continue

        for d2 in os.listdir(p1):
            if d2.startswith('.'):
                continue
            p2 = os.path.join(p1, d2)
            if not os.path.isdir(p2):
                continue

            for d3 in os.listdir(p2):
                if d3.startswith('.'):
                    continue
                p3 = os.path.join(p2, d3)
                if not os.path.isdir(p3):
                    continue

                for d4 in os.listdir(p3):
                    if d4.startswith('.'):
                        continue
                    p4 = os.path.join(p3, d4)
                    if not os.path.isdir(p4):
                        continue

                    for filename in os.listdir(p4):
                        if filename.startswith('.'):
                            continue
                        if filename.endswith("decoder.onnx"):
                            full_path = os.path.join(p4, filename)
                            results.append((d1, d2, d3, d4, filename, full_path,p4 ))

    return results

# Search for checkpoints of models to convert
decoder_entries = search_for_models_decoder(".")

for model_decoder_onnx in decoder_entries:
    print(f"\nExporting to RKNN the decoder ONNX of PiperTTS voice: {model_decoder_onnx[4]}")

    # Load ONNX model and see if the 'g' input exists
    sess = ort.InferenceSession(model_decoder_onnx[5], providers=['CPUExecutionProvider'])

    for inp in sess.get_inputs():
        print(inp.name, inp.shape)

    input_names = [i.name for i in sess.get_inputs()]
    print(input_names)
    input_size_list = [[1, 192, 150], [1, 1, 150]]
    inputs = ['z', 'y_mask']
    if 'g' in input_names:
        input_size_list.append([1, 512, 1])
        inputs.append('g')

    for platform in platforms:
        rknn = RKNN()
        rknn.config(target_platform=platform)
        ret = rknn.load_onnx(model_decoder_onnx[5],
            input_size_list=input_size_list,
            inputs=inputs,
        )
        if ret != 0:
            print('load onnx failed')
            exit(ret)
        ret = rknn.build(do_quantization=False)
        if ret != 0:
            print('build failed')
            exit(ret)
        ret = rknn.export_rknn(f"{model_decoder_onnx[6]}/decoder_{platform}.rknn")
        if ret != 0:
            print('export failed')
            exit(ret)
    
        print(f"Generated decoder_{platform}.rknn in directory: {model_decoder_onnx[6]}\n\n")