# Adapted from: https://github.com/PRBonn/bonnetal/blob/master/train/tasks/segmentation/modules/traceSaver.py
import os
import errno
import onnx
import torch

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


class TraceSaver:
    """
    class to convert the model to ONMX

    :param model_path: Path where pytorch model is saved
    :type model_path: str
    :param new_path: Path to folder where onnx model will be saved as `model.onnx`
    :type new_path:str
    :param force_img_prop:  tuple containing (height, width, channels)
    :type force_img_prop: tuple (1,3)
    """

    def __init__(
        self, model_path: str, new_path: str, force_img_prop: tuple = (None, None, None)
    ):
        """
        :param model_path: path to model config
        :type model_path: str
        :param new_path: new_path
        :type new_path: str
        :param force_img_prop: image transform
        :type force_img_prop: Tuple
        """
        self.path = model_path
        self.new_path = new_path
        if (
            force_img_prop[0] is not None
            and force_img_prop[1] is not None
            and force_img_prop[2] is not None
        ):
            self.height = force_img_prop[0]
            self.width = force_img_prop[1]
            self.channels = force_img_prop[2]
            print("WARNING: FORCING IMAGE PROPERTIES TO")
            print(
                "(Height, Width, Channels: ({}, {}, {}))".format(
                    self.height, self.width, self.channels
                )
            )
        if os.path.exists(self.path):
            self.model = torch.load(self.path)
        else:
            print(f"No trained model found at {self.path}")
            raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), self.path)

        # CUDA speedup
        torch.backends.cudnn.benchmark = True
        self.model = self.model.to(device)

        # eval mode
        self.model.eval()
        print("Creating dummy input to profile")
        self.dummy_input = torch.randn(1, self.channels, self.height, self.width).to(
            device
        )

    def export_onnx(self, debug: bool = True):
        """
        Convert to onnx model with given height and width

        :param debug: print the onnx model graph
        :type debug: bool
        :returns:Save onnx model to folder specified in `new_path` as `model.onnx`

        """
        # create profile
        onnx_path = os.path.join(self.new_path, "model.onnx")
        with torch.no_grad():
            print("Profiling model")
            print("Saving model at ", onnx_path)
            torch.onnx.export(self.model, self.dummy_input, onnx_path)
        # check that it worked
        print("Checking that it all worked out")
        model_onnx = onnx.load(onnx_path)
        onnx.checker.check_model(model_onnx)
        # Print a human readable representation of the graph
        if debug:
            print(onnx.helper.printable_graph(model_onnx.graph))

    def export_dynamic_onnx(self, debug: bool = True):
        """
        Dynamic model onnx used to save a onnx model with different (height, weight) than
        (height, width) used for training model

        :param debug: print the onnx model graph
        :type debug: bool
        :returns:save dynamic onnx model to folder specified in `new_path` as `model_dynamic.onnx`

        """
        # convert to a dynamic ONNX traced model
        # NOTE: Dynamic width/height may not achieve the expected performance improvement
        # with some backend such as TensorRT though.
        input_names = ["input"]
        output_names = ["output"]
        dynamic_axes = {
            "input": {0: "batch_size", 2: "width", 3: "height"},
            "output": {0: "batch_size", 2: "width", 3: "height"},
        }
        # create profile
        onnx_path = os.path.join(self.new_path, "model_dynamic.onnx")
        with torch.no_grad():
            print("Profiling model")
            print("Saving model at ", onnx_path)
            torch.onnx.export(
                self.model,
                self.dummy_input,
                onnx_path,
                input_names,
                output_names,
                dynamic_axes,
            )
        # check that it worked
        print("Checking that it all worked out")
        model_onnx = onnx.load(onnx_path)
        onnx.checker.check_model(model_onnx)
        # Print a human readable representation of the graph
        if debug:
            print(onnx.helper.printable_graph(model_onnx.graph))
