import json
import os
import torch
from fast_scnn import FastSCNN, FastSCNNWithNorm
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"]
MAIN_CLASS_NAME = config["main_class_name"]
N_SHAVES = config["shaves"]
model_to_use = config["model_to_use"]
model_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
model_name = ""
if model_to_use == "geral":
    model_name = f"{MODEL_NAME}_best"
elif model_to_use == "main_class":
    MAIN_CLASS_NAME = config["main_class_name"]
    model_name = f"{MODEL_NAME}_best_f1_{MAIN_CLASS_NAME}"
elif model_to_use == "es":
    ES_CLASSES_NAME = config["es_classes"]
    model_name = f"{MODEL_NAME}_best_es_{ES_CLASSES_NAME}"
else:
    model_name = f"{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)

base = FastSCNNWithNorm(num_classes=NUM_CLASSES, to_rgb=True)  # ajuste num_classes conforme seu labelmap
base.backbone.load_state_dict(torch.load(os.path.join(model_path, f"{model_name}.pth"), map_location="cpu"))
base.eval()

torch.onnx.export(
    base,
    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=N_SHAVES,
    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}")