# indad/indad.py

import os
import sys
from pathlib import Path
import imghdr
from datetime import datetime
import random
import shutil
import json
from typing import Any
from easydict import EasyDict


from PIL import Image
import math

import torch
from torchvision import transforms


sys.path.insert(0, os.path.abspath(
    os.path.join(os.path.dirname(__file__), './../../')
))

import onnxruntime as ort

from ecos_core.indad import indad


FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH


sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


class Model(indad.Model):
    """This is the base model for INDAD

    Args:
        BaseModel (Class): Common class for all ecos_core models

    Returns: None
    """
    NAME = "indad-onnx"
    def __init__(self, opt):
        """This function is called when creating a new instance of INDAD Model

        Args:
            opt (NameSpace): The default options neccessary for initializing the model. 
                             opt should includes these attributes: class_names, weight, device
                             is_pixel_score_used, threshold_score, threshold_pixel_score.

        """
        super().__init__(opt)
        self.opt = opt
        self.opt.class_names = ["good", "anomalous"]
        self.opt.weight = "./model_weight.onnx"
        self.opt.device = "cuda:0" if torch.cuda.is_available() else "cpu"

    def __call__(self, x:torch.tensor):
        """inference with onnx

        Args:
            x (torch.tensor): _description_

        Returns:
            _type_: _description_
        """
        img_level_score, pxl_level_score = self.model.run(['img_level_score', "pxl_level_score"], {'input': x.numpy()})
        return torch.tensor(img_level_score), torch.tensor(pxl_level_score)

    def build_model(self, custom_weight:str=None):
        """Build model in ensemble model.
        Args:
            custom_weight(str): weight to load that not defined in opt. Default: None.
        Returns:
            (InferenceSession): Model after built
        Examples:
            >>> from argparse import ArgumentParser
            >>> from ecos_core.indad.indad_onnx import Model
            >>> parser = ArgumentParser()
            >>> opts = parser.parse_args()
            >>> with open(opt_path, 'r') as f:
            >>>     opts.__dict__ = json.load(f)
            >>> model = Model(opts)
            >>> #Build model with default options
            >>> model.build_model()
            >>> #Build model with custom weight
            >>> model.build_model(custom_weight=<PATH_TO_WEIGHT>)
        """
        
        if custom_weight is None:
            weight_path = os.path.join(str(ROOT), self.opt.weight)
        else:
            weight_path = custom_weight
        self.model = ort.InferenceSession(weight_path)

        return self
    

    