import torch

from blendbatchnorm import fuse_bn_recursively
from meshnet2tfjs import meshnet2tfjs

from meshnet import (
    MeshNet,
    enMesh_checkpoint,
)

device_name = "cuda:0" if torch.cuda.is_available() else "cpu"
device = torch.device(device_name)


def preprocess_image(img, qmin=0.01, qmax=0.99):
    """Unit interval preprocessing"""
    img = (img - img.quantile(qmin)) / (img.quantile(qmax) - img.quantile(qmin))
    return img


# def preprocess_image(img):
#     """Unit interval preprocessing"""
#     img = (img - img.min()) / (img.max() - img.min())
#     return img


# specify how many classes does the model predict
n_classes = 3
# specify the architecture
config_file = "modelAE.json"
# how many channels does the saved model have
model_channels = 15
# path to the saved model
model_path = "model.pth"
# tfjs model output directory
tfjs_model_dir = "model_tfjs"

meshnet_model = enMesh_checkpoint(
    in_channels=1,
    n_classes=n_classes,
    channels=model_channels,
    config_file=config_file,
)

checkpoint = torch.load(model_path)
meshnet_model.load_state_dict(checkpoint)

meshnet_model.eval()

meshnet_model.to(device)
mnm = fuse_bn_recursively(meshnet_model)
del meshnet_model
mnm.model.eval()


meshnet2tfjs(mnm, tfjs_model_dir)