#! /usr/bin/python
from __future__ import division, print_function
import numpy as np
import threading
import time
import io
import PIL
import os
import time
import ujson

# BiSeNetv2
from models.bisenetv2.bisenetv2 import BiSeNetV2
from models.utils import align_and_update_state_dicts
from models.configs.default import cfg

import tensorrt as trt
from torch2trt import torch2trt
from torch2trt import TRTModule

import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

real_color_map = {
                    1 : (255,255,255),    # ice
                    2 : (0, 0, 255),      # water
                    3 : (200,200,200),    # ice-water-mix
                    4 : (42, 150, 42),    # land
                    5 : (200, 200, 255)   # snow-on-land
                    }

sim_color_map = {
                    0 : (0, 0, 0),    # nothing
                    1 : (0, 0, 255),  # blue = frozen water
                    2 : (0, 255, 0),  # green = not frozen water
                    3 : (0, 0, 0)     # trans = conf_threshold
                    }

class ImageDataset(Dataset):
    """
    A torch.Dataset that applies the correct transform to each input image.
    """
    def __init__(self, image_list, transform=None):
        self.image_list = image_list
        self.transform = transform

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, index):
        img = self.image_list[index]

        if self.transform is not None:
            img = np.float32(img).transpose([2,0,1])
            img = self.transform(torch.tensor(img) / 255.0)

        return img.cuda()

class BiSeNetV2Segmentation():
    """
    Acts as a wrapper for the BiSeNetV2 network with utility functions.
    and a __call__ method to use on a standard RGB image.

    :param batch_size: The max size of images batched into the engine.
    :type batch_size: int

    :param weights_path: The filepath to the .ckpt or .pth pytorch file.
    :type weights_path: str

    :param use_trt: Whether or not to attempt a tensorrt conversion of the model.
    :type use_trt: bool

    :param input_w: The width of the input image to resize to.
    :type input_w: int

    :param input_h: The height of the input image to resize to.
    :type input_h: int

    :param num_classes: Number of unique pixel labels for segmentation.
    :type num_classes: int

    :param transform: Pytorch transform to use on the input images.
    :type transform: torch.Transform

    """
    def __init__(self,
                    batch_size,
                    weights_path,
                    use_trt,
                    input_w=3000,
                    input_h=2000,
                    num_classes=5,
                    transform=None,
                    threshold=0
                    ):
        self.use_trt = use_trt
        self.model_name = os.path.basename(weights_path).split('.')[0]
        self.batch_size = batch_size
        self.engine_dir = os.path.dirname(weights_path)
        self.input_w = input_w
        self.input_h = input_h
        self.TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
        self.engine = None
        self.input_names = [ "input" ]
        self.output_names = [ "output" ]
        self.threshold = threshold
        self.threshold_label = 3

        # Set up BiSeNet
        cfg.MODEL.NUM_CLASSES = num_classes
        self.model_lock = threading.RLock()
        os.environ['CUDA_VISIBLE_DEVICES'] = "0"
        use_gpu = torch.cuda.is_available()
        #cfg.merge_from_file('../src/configs/bisenet_noaaice.yaml')
        model = BiSeNetV2(cfg)
        checkpoint = torch.load(weights_path)
        stat_d = checkpoint['state_dict']
        # Converting Lightning Model to Torch model
        for key, val in stat_d.items():
            new_key = key.split('.')[1:]
            new_key = '.'.join(new_key)
            stat_d[new_key] = val
            del stat_d[key]
        print(checkpoint.keys())
        model.load_state_dict(stat_d)
        model.eval().cuda()
        if transform is None:
            transform = transforms.Compose([
                transforms.Normalize(mean=[0.486, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])])
        self.img_transform = transform


        if self.use_trt:
            print("Converting to trt.")
            model = self.create_trt_model( model )

        self.model = model

    def create_trt_model ( self, torch_model ):
        """
        A function to attempt a tensorrt conversion on the given model.
        Loads trt model from .json file if it exists, else it builds
        a trt model after building the pytorch model.

        :param torch_model: Loaded pytorch model.
        :type torch_model: torch.NNModule

        """
        cuda_engine_name = '{}_batchsz{}_{}x{}.engine'.format(
                self.model_name, self.batch_size, self.input_h, self.input_w)
        engine_path = '{}/{}'.format(self.engine_dir, cuda_engine_name)
        io_file_path = engine_path.replace('.engine', '.json')
        if not os.path.exists(engine_path):
            # Build model and load checkpoint weights
            x = torch.randn((1, 3, self.input_h, self.input_w)).cuda()

            print(torch_model)
            print("Converting model to TRT engine.")
            with torch.no_grad():
                _trt_model = torch2trt(torch_model, [x], fp16_mode=True,
                        max_batch_size=self.batch_size)
                engine, input_names, output_names = (_trt_model.engine,
                        _trt_model.input_names, _trt_model.output_names)

            with open(engine_path, 'wb') as f:
                f.write(bytearray(engine.serialize()))

            io_dict = {
                'input_names' : input_names,
                'output_names': output_names
                }

            print('Saving io config to: %s' % io_file_path)

            with open(io_file_path,'w') as jf:
                ujson.dump(io_dict, jf)

        else:
            with open(engine_path, 'rb') as f, trt.Runtime(self.TRT_LOGGER) as runtime:
                engine = runtime.deserialize_cuda_engine(f.read())

                print('Reading io info from {}'.format(io_file_path))

                with open(io_file_path,'r') as pf:
                    io_dict = ujson.load(pf)

                print('Engine loaded')

        model = TRTModule(engine, io_dict['input_names'], io_dict['output_names'])

        print('TRT model built.')
        return model

    def __call__(self, images):
        """
        Calculate pixel labels.

        :param images: List of RGB Images.
        :type images: list of Numpy array

        :return: List of images with labels as pixels.
        :rtype: List of int8 array

        """
        with self.model_lock:
            image_dataset = ImageDataset(images, self.img_transform)
            h, w, d = images[0].shape

            data_loader = DataLoader(image_dataset, batch_size=self.batch_size,
                                     shuffle=False, num_workers=0,
                                     pin_memory=False, drop_last=False)

            labels_list = []
            for batch_idx, imgs in enumerate(data_loader):
                with torch.no_grad():
                    #There are 5 outputs from each stage for deep supervision during train
                    out = self.model(imgs)[0]
                    soft_out = torch.nn.functional.softmax(out[0], dim=0)
                    confs = np.max(soft_out.cpu().numpy())
                    print("confidence % = " + str(confs*100))
                    labels = np.argmax(out.cpu().numpy()[0], axis=0)
                    labels[confs < self.threshold] = self.threshold_label
                    labels_list.append(labels)

            labels = np.concatenate(labels_list, 0)

            return labels
