import json
import os
import torch
from fast_scnn import FastSCNN
from utils import carregar_labelmap_completo

# ⚙️ Configurações
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
RESOLUCAO = config["resolucao"]
model_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
model_name = MODEL_NAME + "_best"

dummy_input = torch.randn(1, 3, RESOLUCAO[1], RESOLUCAO[0])  # (batch, channels, height, width)

_, _, classes, _ = carregar_labelmap_completo(labelmap_path)
NUM_CLASSES = len(classes)

model = FastSCNN(num_classes=NUM_CLASSES)  # ajuste num_classes conforme seu labelmap
model.load_state_dict(torch.load(os.path.join(model_path, model_name + ".pth")))
model.eval()

torch.onnx.export(
    model,
    dummy_input,
    os.path.join(model_path, model_name + ".onnx"),
    input_names=["input"],
    output_names=["output"],
    opset_version=11,
    dynamic_axes=None
)
print(f"Modelo exportado para {os.path.join(model_path, model_name + '.onnx')} com sucesso!")

from openvino.tools.mo import convert_model
from openvino.runtime import serialize
ov_model = convert_model(
    input_model=os.path.join(model_path, model_name + ".onnx"),
    input_shape=[1, 3, RESOLUCAO[1], RESOLUCAO[0]],
    layout="NCHW",
)
serialize(
    ov_model,
    os.path.join(model_path, model_name + ".xml"),
    os.path.join(model_path, model_name + ".bin")
)
print("Conversão para IR concluída e arquivos salvos!")

import blobconverter
blob_path = blobconverter.from_openvino(
    xml=os.path.join(model_path, model_name + ".xml"),
    bin=os.path.join(model_path, model_name + ".bin"),
    data_type="FP16",
    shaves=6,
    output_dir=model_path,
    compile_params=[
        "-ip U8",  # entrada em bytes; compila a conversão interna p/ FP16
        "--mean_values=[123.675,116.28,103.53]",
        "--scale_values=[58.395,57.12,57.375]",
        #"--reverse_input_channels"  # pq você treinou em RGB
    ],
)
print(f"Blob salvo em: {blob_path}")