import shutil
from os import path, scandir, makedirs
from pathlib import Path
import random

import netron
import yaml
from ultralytics import YOLO

from constants import datasets_dir, processed_images_dir


def export_rknn(model_path):
    # model_onnx = export_onnx(model_path)
    model = YOLO(model_path)
    model.export(format="rknn", name="rk3588", simplify=True)
    model_stem = Path(model_path).stem
    folder_name = f'{model_stem}_rknn_model'
    shutil.move(path.join(folder_name, f'{Path(model_path).stem}-rk3588.rknn'), f'{model_stem}.rknn')
    shutil.rmtree(folder_name)
    pass


def export_onnx(model_path):
    model = YOLO(model_path)
    model.export(
        format="onnx",
        imgsz=1280,
        batch=2,
        simplify=True,
        nms=True)
    return Path(model_path).stem + '.onnx'


def export_tensorrt(model_path, dataset_yaml):
    form_data_sample(path.join(path.dirname(dataset_yaml), 'minival', 'images'))
    model = YOLO(model_path)
    with open(dataset_yaml, 'r') as file:
        yaml_data = yaml.safe_load(file) or {}
    yaml_data['minival'] = 'minival/images'
    with open(dataset_yaml, 'w') as file:
        yaml.dump(yaml_data, file)

    model.export(
        format='engine',
        batch=4,
        half=True,
        nms=True,
        data=dataset_yaml,
        split='minival'
    )

def form_data_sample(destination_path, size=500, write_txt_log=False):
    images = []
    with scandir(processed_images_dir) as imd:
        for image_file in imd:
            if not image_file.is_file():
                continue
            images.append(image_file)
    print('shuffling images')
    random.shuffle(images)
    images = images[:size]

    shutil.rmtree(destination_path, ignore_errors=True)
    makedirs(destination_path, exist_ok=True)

    lines = []
    for image in images:
        shutil.copy(image.path, path.join(destination_path, image.name))
        lines.append(f'./{image.name}')
    if write_txt_log:
        with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
            f.writelines([f'{line}\n' for line in lines])

def show_model(model: str = None):
    netron.start(model)


if __name__ == '__main__':
    export_tensorrt('azaion-2025-03-10.pt', path.join(datasets_dir, 'azaion-2025-03-10', 'data.yaml'))
    # export_rknn('azaion-2025-03-10.pt')
    # export_onnx('azaion-2025-03-10.pt')