import pickle
import os
from json import JSONEncoder
from typing import Dict, Optional
import yaml
import json
import numpy as np
from pathlib import Path
from tqdm import tqdm
import math
import torch
from typing import List, Dict, Any
import logging
import coloredlogs
from PIL import Image, ImageDraw
import warnings
import stat
import cv2
import matplotlib.pyplot as plt

def show_anns(anns, borders=True):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:, :, 3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.concatenate([np.random.random(3), [0.5]])
        img[m] = color_mask 
        if borders:
            contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 
            # Try to smooth contours
            contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
            cv2.drawContours(img, contours, -1, (0, 0, 1, 0.4), thickness=1) 

    ax.imshow(img)
    return img

def show_image(image, masks, save_path='demo.png'):
    dpi = 100
    plt.figure(figsize=(image.shape[0]/dpi, image.shape[1]/dpi), dpi=dpi)
    plt.imshow(image)
    mask_img = show_anns(masks)
    plt.axis('off')
    if save_path is None:
        plt.show()
    else:
        plt.close()
        if len(masks) == 0:
            cv2.imwrite(save_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
            return
        new_img = np.uint8(image * (1-mask_img[:,:,3:]) + mask_img[:,:,:3] * mask_img[:,:,3:] *255)
        cv2.imwrite(save_path, cv2.cvtColor(new_img, cv2.COLOR_RGB2BGR))

def save_file(filename,data):
    """
    Based on https://github.com/salesforce/LAVIS/blob/main/lavis/common/utils.py
    Supported:
        .pkl, .pickle, .npy, .json
    """
    parent_dir = os.path.dirname(filename)
    if parent_dir != '':
        if not os.path.exists(parent_dir):
            os.makedirs(parent_dir)
    file_ext = os.path.splitext(filename)[1]
    if file_ext == ".npy":
        with open(filename, "wb+") as fopen:
            np.save(fopen, data)
    elif file_ext == ".json":
        with open(filename,'w+') as fopen:
            json.dump(data,fopen,indent=2)
    elif file_ext == ".yaml":
        with open(filename, "w+") as fopen:
            dump = yaml.dump(data)
            fopen.write(dump)
            fopen.flush()
    else:
        # assume file is pickle
         with open(filename, "wb+") as fopen:
            pickle.dump(data, fopen)

def open_file(filename):
    """
    Based on https://github.com/salesforce/LAVIS/blob/main/lavis/common/utils.py
    Supported:
        .pkl, .pickle, .npy, .json
    """
    file_ext = os.path.splitext(filename)[1]
    if file_ext == '.txt':
        with open(filename,'r+') as fopen:
            data = fopen.readlines()
    elif file_ext in [".npy",".npz"]:
        data = np.load(filename,allow_pickle=True)
    elif file_ext == '.json':
        with open(filename,'r+') as fopen:
            data = json.load(fopen)
    elif file_ext == ".yaml":
        with open(filename,'r+') as fopen:
            data = yaml.load(fopen,Loader=yaml.FullLoader)
    else:
        # assume pickle
        with open(filename,"rb+") as fopen:
            data = pickle.load(fopen)
    return data
