import json
import cv2 as cv
import os
import os.path as osp
from tqdm import tqdm
from PIL import Image
import base64
import numpy as np
from labelme.label_file import LabelFile, PY2, QT4, utils, io
import shutil as sh

__version__ = "0.1.0"

def encodeImage(imarr, ext=".jpg"):
    """@See Also LabelFile.load_image_file

    Arguments:
        imarr(np.ndarray | PIL.Image | str)

    Return:
        (str) base64 encoded image
    """
    imtype = type(imarr)
    if isinstance(imarr, Image.Image):
        pass
    elif imtype == np.ndarray:
        imarr = Image.fromarray(imarr)
    elif imtype == str:
        imarr = PIL.Image.open(imarr)
    else:
        raise ValueError("Unsupported Type: ", imtype)
    # apply orientation to image according to exif
    image_pil = utils.apply_exif_orientation(imarr)

    with io.BytesIO() as f:
        # ext = osp.splitext(filename)[1].lower()
        if PY2 and QT4:
            format = "PNG"
        elif ext in [".jpg", ".jpeg"]:
            format = "JPEG"
        else:
            format = "PNG"
        image_pil.save(f, format=format)
        f.seek(0)
        # return f.read()
        imageData = f.read()
    return base64.b64encode(imageData).decode("utf-8")


def calcY(line, x):
    return line[0] * x + line[1]

import copy
class LabelmeModifier(object):
    """
    Arguments:
        root(str): root directory that contains the raw image and json file
        generated by `labelme`
        dst_dir(str)
    """

    def __init__(self, root, dst_dir, image_suffix=".jpg", debug=False, **kwargs):
        self.root = root
        self.dst_dir = dst_dir
        self.debug = debug
        self.image_suffix = image_suffix
        files = [f for f in sorted(os.listdir(root))]
        self.images = [file for file in files if image_suffix in file]
        jsonfiles = [file for file in files if ".json" in file ]
        self.names = [img[: img.rfind(image_suffix)] for img in self.images]

        print(f"Found files at {root} images: {len(self.images)}, jsons: {len(jsonfiles)}")
        
        os.makedirs(self.dst_dir, exist_ok=True)


    def readJsonFile(self, file):
        loc = osp.join(self.root, file)
        if not osp.exists(loc):
            return None
        try:
            with open(loc, "r") as f:
                content = json.loads(f.read())
        except Exception as e:
            print(e)
            raise ValueError(f"Exception for {file}: ")
        return content


    def _resize(self, name, size, jsoncontent):
        """
        Note: **imageData in jsoncontent is not modified, it should be modified by the last step.**

        Arguments:
            name(str): filename stem without path location
            size(tuple): new size that resized image will be
            jsoncontent: correspoinding labelme json file content

        Return:
            tuple: (newImage, jsoncontent)
        """
        # resize images only
        image = Image.open(osp.join(self.root, f"{name}{self.image_suffix}"))
        iw, ih = image.size
        wscale = size[0] / iw
        hscale = size[1] / ih
        newImage = image.resize(size)

        if jsoncontent is None:
            if self.debug:
                print(f"Not Found json file for {name}")
            return newImage, None
        jsoncontent["imageWidth"]  = size[0]
        jsoncontent["imageHeight"] = size[1]
        # if self.debug:
            # newImage.show(f"Resized Image {name}")
        shapes = jsoncontent["shapes"]
        for shape in shapes:
            for point in shape["points"]:
                point[0] *= wscale
                point[1] *= hscale
        
        if self.debug:
            print("Name: ", name)
        return newImage, jsoncontent

    def resize(self, size):
        """
        Collabrate with crop, If crop is need before resize, using cropResize instead

        Arguemengs:
            size(tuple): W X H
        """
        count = 0
        for name in tqdm(self.names):
            jsoncontent = self.readJsonFile(f"{name}.json")
            # print("JsonContent", jsoncontent)
            newImage, jsoncontent = self._resize(name, size, jsoncontent)
            savedLoc = osp.join(self.dst_dir, f"{name}{self.image_suffix}")
            newImage.save(savedLoc)
            imageData = LabelFile.load_image_file(savedLoc)
            imageData = base64.b64encode(imageData).decode("utf-8")
            if jsoncontent is not None:
                jsoncontent["imageData"] = imageData

                with open(osp.join(self.dst_dir, f"{name}.json"), "w") as f:
                    f.write(json.dumps(jsoncontent, indent=True))
            count += 1
            if self.debug:
                break

        print(f"Done for {count}")

    def _rotate(self, name, jsoncontent, degree, center, boarderValue, scale=1.0):
        """
        Arguments: 
            name(str): Image Name
            degree(float): degrees
            center(tuple): X,Y

        See torchutils.rotateImage
        """
        image = Image.open(osp.join(self.root, f"{name}{self.image_suffix}"))
        iw, ih = image.size    #    
        if center is None:
            center = (iw // 2, ih // 2 )
        image = np.asarray(image)
        M = cv.getRotationMatrix2D(center, degree, scale)
        rotated = cv.warpAffine(image, M, (iw, ih), borderValue=boarderValue)
        # if self.debug:
        #     cv.imshow("default", rotated)
        #     cv.waitKey(0)

        if jsoncontent is None:
            return Image.fromarray(rotated)

        shapes = jsoncontent["shapes"]
        for shape in shapes:
            points = []
            # rotatepoints = cv.warpAffine(points, M, (iw, ih), borderValue=boarderValue)
            # print(rotatepoints.shape, points.shape)
            for point in shape["points"]:
                newPoint = np.dot(M[:,:2], np.asarray(point)) + M[:,2]
                # print("\n=======", newPoint, M[:, 2])
                # check point pos
                if newPoint[0] > iw: newPoint[0] = iw
                if newPoint[1] > ih: newPoint[0] = ih
                points.append([newPoint[0], newPoint[1]])
                # if self.debug: break
            # if self.debug: break
            shape["points"] = points
        return Image.fromarray(rotated)
        

    def rotate(self, start=-10, end=10, num=20, center=None, boarderValue = (0, 0, 0)):
        """
        Arguments:
            start(float): start_degree for rotation
            end(float): end_degree for rotation
            interval(float): interval_degree for rotation
            center(tuple | None): if the all the image size are the same, it could be applied when rotating,
                if image size differes from each other, leave it unchanged to let program choose the center 
                of the image
        """
        count = 0
        for name in tqdm(self.names):
            orijsoncontent = self.readJsonFile(f"{name}.json")
            if orijsoncontent is None:
                continue
            # print("JsonContent", jsoncontent)
            rc = 0
            for deg in np.linspace(start, end, num):
                jsoncontent = copy.deepcopy(orijsoncontent) 
                newImage = self._rotate(name, jsoncontent, deg, center, boarderValue)
                newName = f"{name}-r{rc:04}"
                savedLoc = osp.join(self.dst_dir, f"{newName}{self.image_suffix}")
                rc += 1
                newImage.save(savedLoc)
                imageData = LabelFile.load_image_file(savedLoc)
                imageData = base64.b64encode(imageData).decode("utf-8")
                if jsoncontent is not None:
                    jsoncontent["imageData"] = imageData

                    with open(osp.join(self.dst_dir, f"{newName}.json"), "w") as f:
                        f.write(json.dumps(jsoncontent, indent=True))
                count += 1

                if self.debug:
                    break
            if self.debug:
                break

        print(f"Done for {count}")

    def cropResize(self, leftTop, cropSize, newSize = None):
        """
        Arguements:
            leftTop(tuple): (X, Y)

            cropSize(tuple): W X H, croped size.  
                (cropSize + leftTop) should be less than original_size at W & H, 
                if not, (cropSize + leftTop) will reach the original images rightBottom

            newSize(tuple): W X H, rescaled size after cropped, if it's None, then it 
                won't be resized
        """
        raise NotImplementedError("Not Impletmented Yet")
        

    def output(self, destdir):
        """
        Arguments:
            destdir(str): dest directory
        """
        pass

    def removeLabels(self, labels):
        """
        Arugments:
            labels(list(str))
        """
        assert type(labels) is list
        slen = len(labels)
        si = [False for _ in range(slen)]
        nr = [0 for _ in range(slen)]
        self._renameShapes(labels, None, shape_idxes=si, num_reserved=nr)

    def _renameShapes(self, shapenames, new_shapenames=None, shape_idxes=None, num_reserved=None):
        """Rename one shape at one time,

        Arguments:
            shapenames(list(str)): shape that should be renamed or decreased
            new_shapenames(list(str)): if set, old shapenames will be renamed
            num_reserved(list(int)): if set, `num_reserved` shapes will be reversed, if it equal to or less than 0, then specified shape will be deleted.
            shape_idxes(list(bool)): if True, new shapenames will be concatenated with `-1` or `-2` suffix
        """
        count = 0
        if new_shapenames is None:
            new_shapenames = shapenames

        assert len(new_shapenames) == len(shapenames) == len(shape_idxes) == len(num_reserved)

        copyImage = (self.root != self.dst_dir)

        for name in tqdm(self.names):
            orijsoncontent = self.readJsonFile(f"{name}.json")
            if orijsoncontent is None:
                continue
            new_shapes = []
            reserve_count = 0
            # iterate over original json content
            for shape in orijsoncontent["shapes"]:
                if shape["label"] in shapenames:
                    idx = shapenames.index(shape["label"])
                    # if count not meet, reserve it
                    if reserve_count < num_reserved[idx]:
                        nshape = copy.deepcopy(shape)
                        if shape_idxes[idx]:
                            nshape["label"] = new_shapenames[idx] + f"-{reserve_count}"
                        else:
                            nshape["label"] = new_shapenames[idx]
                        new_shapes.append(nshape)
                        reserve_count += 1
                else:
                    new_shapes.append(shape)
            # print(new_shapes)
            orijsoncontent["shapes"] = list(sorted(new_shapes, key=lambda x: x['label']))
            with open(osp.join(self.dst_dir, f"{name}.json"), "w") as f:
                f.write(json.dumps(orijsoncontent, indent=True))
            if copyImage:
                sh.copy(osp.join(self.root, f"{name}{self.image_suffix}"), self.dst_dir)
            count += 1
        print(f"Done for {count}")

    def renameShapes(self, shapename, new_shapename=None, num_reserved=None, shape_idxes=None):
        """
        Arguments:
            shapename (list(str) | str)
            new_shapename (list(str) | str)
        """
        stype = type(shapename)
        if stype == str:
            shapename = [shapename]
        assert type(shapename) is list
        slen = len(shapename)
        if shape_idxes is None:
            shape_idxes = [False for _ in range(slen)]
        if num_reserved is None:
            num_reserved = [1000 for _ in range(slen)]
        self._renameShapes(shapename, new_shapename, shape_idxes=shape_idxes, num_reserved=num_reserved)

        
    def polyfitSlope(self):
        for jsonfile in tqdm(jsonfiles):
        #     print(jsonfile)
            with open(src/jsonfile, "r") as f:
                content = json.loads(f.read())
            # print(content)
            shapes = content['shapes']
            imagePath = content['imagePath']
            occipitalSlope = None
            sphenoidbone = None
            newShapes = []
            for i, shape in enumerate(shapes):
                if shape['label'] == 'SphenoidBone':
                    sphenoidbone = shape
                elif shape['label'] == labelType:
                    occipitalSlope = shape
                    continue
                newShapes.append(shape)
                
            shapes = newShapes
                    
        #     if occipitalSlope is not None:
        #         continue
            imdata = cv.imread(str(src/ imagePath), 0)
        #     flatcontour = np.asarray(sphenoidbone["points"], dtype=int)
        #     contourLen = len(flatcontour)
            # print(flatcontour)
        #     maxy = (np.argmax(flatcontour[:, 1]))
        #     maxy = (maxy + 1) % len(flatcontour)
        #     if (maxy + 5) < len(flatcontour):
        #         contour = flatcontour[maxy:maxy+5]
        #     else:
        #         contour = flatcontour[maxy:]
        #         con = contour.shape[0]
        # #         print("length", con)
        #         contour = np.concatenate((contour, flatcontour[: 5-con]))
        #     line = np.polyfit(contour[:,0].reshape(-1), contour[:,1].reshape(-1), 1)
        #     x1 = float(contour[0, 0])
        #     x2 = float(contour[-1, 0] + 20)
        #     y1 = float(calcY(line, x1))
        #     y2 = float(calcY(line, x2))
        #     print([[x1, y1], [x2, y2]])
            # cv.line(imdata, (x1, y1), (x2, y2),(0,255,0), 2) 
            # plt.figure(figsize=(10, 8 ), dpi=100)
            # plt.imshow(imdata, cmap='gray', vmin=0, vmax=255)
        #     occipitalSlope = {"label": labelType, "points": [[x1, y1], [x2, y2], [x2, y2+1], [x1, y1+1]], "gorup_id": None, "shape_type": "polygon", "flags": {}}
        #     shapes.append(occipitalSlope)
        #     print([s["label"] for s in shapes], len(shapes))
        #     break
            content["shapes"] = shapes
            try:
                data = json.dumps(content, indent=True)
                
                with open(src/jsonfile, "w") as f:
                    f.write(data)

            except Exception as e:
                    print(e)


def moveImages(src, dst):
    """move Images from src dir with json files not moved
    """
    jsons = []
    imgs = []
    for root, dirs, files in os.walk(src):
        for file in files:
            if file.endswith(".json"):
                jsons.append(file[:file.rfind(".")])
            else:
                imgs.append(file)
    
    dst = str(dst)
    for file in imgs:
        filename = file[:file.rfind(".")]
        if filename in jsons:
            continue
        # print(filename)
        sh.move(str(src / file), dst)


import argparse
import glob
import os
import os.path as osp
import sys

import imgviz
import numpy as np

import labelme

class Labelme2Vocor():

    def __init__(self, input_dir, output_dir, labels, noviz=False, debug=False):
        if osp.exists(output_dir):
            print("Output directory already exists: ", output_dir, " Deleting...")
            import shutil as sh
            sh.rmtree(output_dir)
        os.makedirs(output_dir)
        os.makedirs(osp.join(output_dir, "JPEGImages"))
        os.makedirs(osp.join(output_dir, "SegmentationClass"))
        os.makedirs(osp.join(output_dir, "SegmentationClassPNG"))

        if not noviz:
            os.makedirs(osp.join(output_dir, "SegmentationClassVisualization"))
        print("Creating dataset:", output_dir)

        self.input_dir = input_dir
        self.output_dir = output_dir
        self.labels = labels
        self.noviz = noviz
        self.debug = debug

    def getClasses(self):
        ## get class names
        class_names = []
        class_name_to_id = {}
        for i, line in enumerate(open(self.labels).readlines()):
            class_id = i - 1  # starts with -1
            class_name = line.strip()
            if len(class_name) == 0:
                continue
            class_name_to_id[class_name] = class_id
            if class_id == -1:
                assert class_name == "__ignore__"
                continue
            elif class_id == 0:
                assert class_name == "_background_"
            class_names.append(class_name)

        class_names = tuple(class_names)
        self.class_names = class_names
        self.class_name_to_id = class_name_to_id
        print("class_names:", class_names)
        out_class_names_file = osp.join(self.output_dir, "class_names.txt")

        with open(out_class_names_file, "w") as f:
            f.writelines("\n".join(class_names))
        print("Saved class_names:", out_class_names_file)

    def output(self, ordered_keys=None):
        """
        Arguments:
            ordered_keys(list(str))
        """
        for filename in tqdm(glob.glob(osp.join(self.input_dir, "*.json"))):
            # print("Generating dataset from:", filename)

            label_file = labelme.LabelFile(filename=filename)

            base = osp.splitext(osp.basename(filename))[0]
            out_img_file = osp.join(self.output_dir, "JPEGImages", base + ".jpg")
            out_lbl_file = osp.join(
                self.output_dir, "SegmentationClass", base + ".npy"
            )
            out_png_file = osp.join(
                self.output_dir, "SegmentationClassPNG", base + ".png"
            )
            if not self.noviz:
                out_viz_file = osp.join(
                    self.output_dir,
                    "SegmentationClassVisualization",
                    base + ".jpg",
                )

            with open(out_img_file, "wb") as f:
                f.write(label_file.imageData)
            img = labelme.utils.img_data_to_arr(label_file.imageData)
            if ordered_keys is not None:
                newshapes = []
                for ok in ordered_keys:
                    for shape in label_file.shapes:
                        if shape["label"] == ok:
                            newshapes.append(shape)
                label_file.shapes=newshapes

            if self.debug:
                print(label_file.shapes)
                break

            lbl, _ = labelme.utils.shapes_to_label(
                img_shape=img.shape,
                shapes=label_file.shapes,
                label_name_to_value=self.class_name_to_id,
            )
            labelme.utils.lblsave(out_png_file, lbl)

            # np.save(out_lbl_file, lbl)
                
            if not self.noviz:
                if img.shape[0] == 1: # gray img
                    img = imgviz.rgb2gray(img)
                viz = imgviz.label2rgb(
                    label=lbl,
                    # img=imgviz.rgb2gray(img),
                    img=img,
                    font_size=15,
                    label_names=self.class_names,
                    loc="rb",
                )
                imgviz.io.imsave(out_viz_file, viz)



import collections
import datetime
import glob
import uuid

import numpy as np
import labelme
import re
SubCategoryPatter = re.compile("(\w+)-(\w+|\d+)")

class Labelme2Cocoer():
    def __init__(self, output_dir, cocomask=None, debug=False):
        """
        Arguments:
            output_dir(str)
            cocomask(pytcocotools.mask): if none, default pycocotools.mask will be used
        """
        self.output_dir = output_dir
        self.debug = debug

        if cocomask is None:
            try:
                import pycocotools.mask
            except ImportError:
                print("Please install pycocotools:\n\n    pip install pycocotools\n")
                sys.exit(1)
            cocomask = pycocotools.mask

        self.cocomask = cocomask

        now = datetime.datetime.now()
        self.data = dict(
            info=dict(
                description=None,
                url=None,
                version=None,
                year=now.year,
                contributor=None,
                date_created=now.strftime("%Y-%m-%d %H:%M:%S.%f"),
            ),
            licenses=[dict(url=None, id=0, name=None,)],
            images=[
                # license, url, file_name, height, width, date_captured, id
            ],
            type="instances",
            annotations=[
                # segmentation, area, iscrowd, image_id, bbox, category_id, id
            ],
            # supercategory, id, name
            categories= None,
        )

    def classNameToId(self, labels_file, skeletondict=None):
        """
        Currently keypoint labelname should be `xxx-1` or `xxx-2`, because it's
        easy to be ordered.

        Argument:
            labels_file (str): labels file loc,
        """
        class_id = -1
        # dict to map class_name (str) to id (int)
        class_name_to_id = {}
        categoriedict = {}
        with open(labels_file, "r") as f:
            for i, line in enumerate(f.readlines()):
                class_name = line.strip()   
                if len(class_name) == 0:
                    continue
                
                if class_id == -1:
                    assert class_name == "__ignore__"
                    class_id += 1
                    continue
                matched = SubCategoryPatter.match(class_name)
                if matched is None:
                    class_name_to_id[class_name] = class_id
                    categoriedict[class_id] = { 
                        "name": class_name, 
                        "supercategory": None,
                        "type": "object", 
                        "keypoints": [],
                    }
                    class_id += 1
                else:
                    cid = class_name_to_id[matched[1]]
                    class_name_to_id[class_name] = cid
                    cat = categoriedict[cid]
                    cat["supercategory"] = None
                    cat["type"] = "keypoint"
                    cat["keypoints"].append(class_name)
        if self.debug:
            print(class_name_to_id, "============")
            # raise ValueError("Test")
        self.class_name_to_id = class_name_to_id
        self.categoriedict = categoriedict

        categories = []
        kp_cats = {}
        skeleton = [] if skeletondict is None else skeletondict
        for class_id, cat in self.categoriedict.items():
            if cat["type"] == "object":
                categories.append({
                    "supercategory": cat["supercategory"], 
                    "id": class_id, 
                    "name": cat["name"],
                })
            elif cat["type"] == "keypoint":
                categories.append({
                    "supercategory": cat["supercategory"], 
                    "id": class_id, 
                    "name": cat["name"],
                    "keypoints": cat["keypoints"],
                    "skeleton": skeleton
                })

        self.data["categories"] = list(sorted(categories, key=lambda x: x["id"]))

        if self.debug:
            print(self.data["categories"], "============\n\n")

    def _getAnno(self, label_file, img, image_id):
        """
        Extract all annotations from single labeled json file.

        Arguments:
            label_file: label_file object
            img: img_arr
            image_id (int): 
        """
        masks = {}  # for area
        segmentations = collections.defaultdict(list)  # for segmentation
        keypoints = {}
        for shape in label_file.shapes:
            shape_type = shape.get("shape_type", "polygon")

            points = shape["points"]
            label = shape["label"]
            group_id = shape.get("group_id")
            if shape_type == "point": 
                # keypoint  is special
                matched = SubCategoryPatter.match(label)
                if matched[1] not in keypoints:
                    keypoints[matched[1]] = []
                keypoints[matched[1]].append((matched[2], label, points[0]))
                continue

            mask = labelme.utils.shape_to_mask(img.shape[:2], points, shape_type)
            if group_id is None:
                group_id = uuid.uuid1()

            instance = (label, group_id)

            # if some instance is the same, then their masks should be intersected
            if instance in masks:
                masks[instance] = masks[instance] | mask
            else:
                masks[instance] = mask

            if shape_type == "rectangle":
                (x1, y1), (x2, y2) = points
                x1, x2 = sorted([x1, x2])
                y1, y2 = sorted([y1, y2])
                # clock wise
                points = [x1, y1, x2, y1, x2, y2, x1, y2]
            else:
                points = np.asarray(points).flatten().tolist()

            segmentations[instance].append(points)
        # sort keypoints to get ordered value
        for _, kp in keypoints.items():
            kp.sort(key=lambda x: x[0])

        segmentations = dict(segmentations)
        if self.debug:
            print("Mask: ", len(masks), masks.keys())
            print("Segm", len(segmentations))
            # print(" \n\n ")
            # print(segmentations)
            # print(" \n\n ")
            # print(keypoints)
            self.masks = masks

        annotations = {}
        for instance, mask in masks.items():
            cls_name, group_id = instance
            if cls_name not in self.class_name_to_id:
                continue
            cls_id = self.class_name_to_id[cls_name]

            mask = np.asfortranarray(mask.astype(np.uint8))
            mask = self.cocomask.encode(mask)
            bbox = self.cocomask.toBbox(mask).flatten().tolist()
            # it calculates mask area, not box area that torch needs
            area = float(self.cocomask.area(mask)) 
            # following is the area that torch needs
            # area = float(bbox[2] * bbox[3])
            # bbox = [bbox[0], bbox[1], bbox[0], bbox[1]] # format for torch

            anno = dict(
                    id=len(self.data["annotations"]),
                    image_id=image_id,
                    category_id=cls_id,
                    segmentation=segmentations[instance],
                    area=area,
                    bbox=bbox,
                    iscrowd=0,
                )
            if cls_name in keypoints:
                kps = keypoints[cls_name] # 3-element-tuple list
                kplist = []
                num = 0
                # print(kps, "========")
                for kp in kps:
                    # kp is the 3-element-tuple list
                    k = kp[2]
                    kplist.append(k[0])
                    kplist.append(k[1])
                    kplist.append(2)
                    num += 1
                anno.update({
                    "keypoints": kplist, # [[kp[0], kp[1], 1] for kp in kps ], # format for torch
                    "num_keypoints": num,
                })

            # annotations.append(anno)
            self.data["annotations"].append(anno)

        if self.debug:
            print(self.data["annotations"])
            # raise ValueError("Test")

    def generateCocoJson(self, label_files):
        """
        label_files: iterateble files
        """
        out_ann_file = osp.join(self.output_dir, "annotations.json")
        for image_id, filename in tqdm(enumerate(label_files)):
            # print("Generating dataset from:", filename)

            label_file = labelme.LabelFile(filename=filename)

            base = osp.splitext(osp.basename(filename))[0]
            out_img_file = osp.join(self.output_dir, "JPEGImages", base + ".jpg")

            img = labelme.utils.img_data_to_arr(label_file.imageData)
            Image.fromarray(img).convert("RGB").save(out_img_file)
            self.data["images"].append(
                dict(
                    license=0,
                    url=None,
                    file_name=osp.relpath(out_img_file, osp.dirname(out_ann_file)),
                    height=img.shape[0],
                    width=img.shape[1],
                    date_captured=None,
                    id=image_id,
                )
            )

            self._getAnno(label_file, img, image_id)

            if self.debug:
                break


    def output(self, out_ann_file = None):
        if out_ann_file is None:
            out_ann_file = "annotations.json"
        out_ann_file = osp.join(self.output_dir, out_ann_file)
        with open(out_ann_file, "w") as f:
            json.dump(self.data, f, indent=True)
