from PIL import Image
import numpy as np
from utils import AppUtils
import cv2
from inference import start_inference

class AppInference:
    def __init__(self):
        self.COLOR_MAP = {}    

    def inference(self, input_id, img_path, mask, label):
        AppUtils.clear()
        self._input_id = input_id
        self._handle_preprocess(img_path, mask, label)
        return self._handle_model_inference()
    
    def _handle_preprocess(self, img_path, mask, label):
        items = self._read_files(img_path)
        self._items = items
        mask = self._save_mask(items, mask)
        if label != "None":
            self._edit_maps(items, mask, label, save=True)

    def _handle_model_inference(self):
        return start_inference()
    
    def preview(self, input_id, img_path, mask, label):
        AppUtils.clear()
        self._input_id = input_id
        items = self._read_files(img_path)
        mask = self._save_mask(items, mask)
        if label != "None":
            self._edit_maps(items, mask, label)
        return self.generate_colored_image(items["inst_map"])

    def generate_colored_image(self, semantic_map):
        np.random.seed(256)
        if len(semantic_map.shape) == 3:
            semantic_map = semantic_map[:,:,1]
        color_image = np.zeros((semantic_map.shape[0], semantic_map.shape[1], 3), dtype=np.uint8)
        for row in range(semantic_map.shape[0]):
            for col in range(semantic_map.shape[1]):
                inst_id = semantic_map[row, col]
                if self.COLOR_MAP.get(inst_id, None) is None:
                    self.COLOR_MAP[inst_id] = np.random.randint(256, size=(3,))
                color_image[row, col, :] = self.COLOR_MAP[inst_id]
        return Image.fromarray(color_image)

    def _read_files(self, img_path):
        dataset = img_path.split("/")[2]
        items = {
            "img_path": img_path,
            "label_path": img_path.replace("images", "labels").replace("jpg", "png"),
            "inst_map_path": img_path.replace("images", "inst_map").replace("jpg", "png"),
        }
        for file_path in items.values():
            AppUtils.copy_file(file_path, file_path.replace(dataset, "test_processed"))
        items["dataset"] = dataset
        base_img = cv2.imread(img_path)
        base_img = cv2.cvtColor(base_img, cv2.COLOR_BGR2RGB)
        base_lab = cv2.imread(items["label_path"], 0)
        base_inst_map = Image.open(items["inst_map_path"])            
        base_inst_map = np.array(base_inst_map, dtype=np.int32)
        items.update(
            {
                "img": base_img,
                "label": base_lab,
                "inst_map": base_inst_map,
            }
        )
        return items

    def _save_mask(self, items, mask):
        mask = np.array(mask)[:,:,0]
        mask = mask.reshape((1,) + mask.shape).astype(np.float32)
        save_path = items["img_path"].replace(items["dataset"], "test_processed").replace("images", "predefined_masks/type_0").replace("jpg", "png")
        cv2.imwrite(save_path, mask[0]* 255)
        return mask[0].astype(np.uint8)

    def _edit_maps(self, items, mask, label, save=False):
        mask_path = items["img_path"].replace(items["dataset"], "test_processed").replace("images", "predefined_masks/type_0").replace("jpg", "png")
        mask = cv2.imread(mask_path, 0) / 255
        target_pixels = mask == 1
        target_inst_id = AppUtils.get_inst_id(self._input_id, label)
        items["inst_map"][target_pixels] = target_inst_id
        items["label"][target_pixels] = (target_inst_id % 120)
        im = Image.fromarray(items["inst_map"]).convert("I")
        if save:
            im.save(items["inst_map_path"].replace(items["dataset"], "test_processed"))
            cv2.imwrite(items["label_path"].replace(items["dataset"], "test_processed"), items["label"])