import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import torch

from annotator.oneformer.detectron2.config import get_cfg
from annotator.oneformer.detectron2.projects.deeplab import add_deeplab_config
from annotator.oneformer.detectron2.data import MetadataCatalog

from annotator.oneformer.oneformer import (
    add_oneformer_config,
    add_common_config,
    add_swin_config,
    add_dinat_config,
)

from annotator.oneformer.oneformer.demo.defaults import DefaultPredictor
from annotator.oneformer.oneformer.demo.visualizer import Visualizer, ColorMode


def make_detectron2_model(config_path, ckpt_path):
    cfg = get_cfg()
    add_deeplab_config(cfg)
    add_common_config(cfg)
    add_swin_config(cfg)
    add_oneformer_config(cfg)
    add_dinat_config(cfg)
    cfg.merge_from_file(config_path)
    if torch.cuda.is_available():
        cfg.MODEL.DEVICE = 'cuda'
    else:
        cfg.MODEL.DEVICE = 'cpu'
    cfg.MODEL.WEIGHTS = ckpt_path
    cfg.freeze()
    metadata = MetadataCatalog.get(cfg.DATASETS.TEST_PANOPTIC[0] if len(cfg.DATASETS.TEST_PANOPTIC) else "__unused")
    return DefaultPredictor(cfg), metadata


def semantic_run(img, predictor, metadata):
    predictions = predictor(img[:, :, ::-1], "semantic")  # Predictor of OneFormer must use BGR image !!!
    visualizer_map = Visualizer(img, is_img=False, metadata=metadata, instance_mode=ColorMode.IMAGE)
    out_map = visualizer_map.draw_sem_seg(predictions["sem_seg"].argmax(dim=0).cpu(), alpha=1, is_text=False).get_image()
    return out_map