import json 
import numpy as np

from matplotlib import cm
import matplotlib
from PIL import Image, ImageColor, ImageFont, ImageDraw 
import numpy as np
import pdb
from datetime import date
today = date.today()
FONTS = {'amiko': "fonts/Amiko-Regular.ttf",
        'nature': "fonts/LoveNature.otf", 
        'painter':"fonts/PainterDecorator.otf",
        'animals': "fonts/UncialAnimals.ttf", 
        'zen': "fonts/ZEN.TTF"}

#########################################
#  Draw keypoints on image
def draw_keypoints_on_image(image,
                            keypoints,
                            map_label_id_to_str,
                            flag_show_str_labels,
                            use_normalized_coordinates=True,                           
                            font_style='amiko',
                            font_size=8,
                            keypt_color="#ff0000",
                            marker_size=2,
                            color_by_confidence=True,
                            ):
    """Draws keypoints on an image.
    Modified from:
        https://www.programcreek.com/python/?code=fjchange%2Fobject_centric_VAD%2Fobject_centric_VAD-master%2Fobject_detection%2Futils%2Fvisualization_utils.py
    Args:
    image: a PIL.Image object.
    keypoints: a numpy array with shape [num_keypoints, 2].
    map_label_id_to_str: dict with keys=label number and values= label string
    flag_show_str_labels: boolean to select whether or not to show string labels
    color: color to draw the keypoints with. Default is red.
    radius: keypoint radius. Default value is 2.
    use_normalized_coordinates: if True (default), treat keypoint values as
        relative to the image.  Otherwise treat them as absolute.

    
    """
    # get a drawing context
    draw = ImageDraw.Draw(image,"RGBA")  

    im_width, im_height = image.size
    keypoints_x = [k[0] for k in keypoints]
    keypoints_y = [k[1] for k in keypoints]
    confidences = [k[2] for k in keypoints]

    # adjust keypoints coords if required
    if use_normalized_coordinates:
        keypoints_x = tuple([im_width * x for x in keypoints_x])
        keypoints_y = tuple([im_height * y for y in keypoints_y])
    
    #cmap = matplotlib.cm.get_cmap('hsv')
    # draw ellipses around keypoints
    for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
        # handling potential nans in the keypoints
        if np.isnan(keypoint_x).any():
            continue

        confidence = float(np.clip(confidences[i], 0, 1))
        if color_by_confidence:
            # fill color encodes the keypoint confidence (see add_confidence_legend)
            round_fill = cm.viridis(confidence, bytes=True)
        else:
            # one color per bodypart, transparency encodes the confidence
            round_fill = list(cm.viridis(i / max(len(keypoints) - 1, 1), bytes=True))
            round_fill[3] = round(confidence * 255)
            round_fill = tuple(round_fill)
        draw.ellipse([(keypoint_x - marker_size, keypoint_y - marker_size),
                      (keypoint_x + marker_size, keypoint_y + marker_size)],
                       fill=tuple(round_fill), outline= 'black', width=1) #fill and outline: [0,255]

        # add string labels around keypoints
        if flag_show_str_labels:
            font = ImageFont.truetype(FONTS[font_style], 
                                     font_size) 
            draw.text((keypoint_x + marker_size, keypoint_y + marker_size),#(0.5*im_width, 0.5*im_height), #-------
                      map_label_id_to_str[i],
                      ImageColor.getcolor(keypt_color, "RGB"), # rgb #
                      font=font)

#########################################
#  Legend for the keypoint confidence colors
def add_confidence_legend(image, font_style='amiko'):
    """Returns the image with a white band below it holding a viridis strip (confidence 0 to 1).

    The band is added below the image, so the legend never covers it and keypoint
    coordinates are unchanged.
    """
    im_width, im_height = image.size
    strip_w = max(60, im_width // 5)
    strip_h = max(6, im_height // 60)
    font = ImageFont.truetype(FONTS[font_style], max(10, strip_h * 2))
    margin = strip_h
    band_h = strip_h + font.size + 3 * margin

    out = Image.new("RGB", (im_width, im_height + band_h), "white")
    out.paste(image.convert("RGB"), (0, 0))
    draw = ImageDraw.Draw(out)
    x0 = im_width - strip_w - 2 * margin
    y0 = im_height + margin
    for dx in range(strip_w):
        draw.line([(x0 + dx, y0), (x0 + dx, y0 + strip_h)],
                  fill=cm.viridis(dx / (strip_w - 1), bytes=True))
    label_y = y0 + strip_h + margin // 2
    draw.text((x0, label_y), "0", fill="black", font=font)
    draw.text((x0 + strip_w, label_y), "1", fill="black", font=font, anchor="ra")
    draw.text((x0 + strip_w / 2, label_y), "confidence", fill="black", font=font, anchor="ma")
    return out


#########################################
#  Keypoint confidences as table rows
def keypoint_confidence_rows(kpts_per_animal, map_label_id_to_str):
    """(animal, bodypart, confidence) for every keypoint kept (not NaN), lowest confidence first."""
    rows = []
    for i_animal, kpts in enumerate(kpts_per_animal):
        for i_kpt, kpt in enumerate(kpts):
            if not np.isnan(kpt[2]):
                rows.append([i_animal, map_label_id_to_str[i_kpt], round(float(kpt[2]), 3)])
    return sorted(rows, key=lambda row: row[2])


#########################################
#  Save the annotated image for download
def save_annotated_image(image, path_to_output_file='download_annotated.png'):
    image.save(path_to_output_file)
    return path_to_output_file


#########################################
#  Draw bboxes on image
def draw_bbox_w_text(img,
                     results,
                     font_style='amiko',
                     font_size=8): #TODO: select color too?
    #pdb.set_trace()
    bbxyxy = results
    w, h = bbxyxy[2], bbxyxy[3]
    shape = [(bbxyxy[0], bbxyxy[1]), (w , h)]
    imgR = ImageDraw.Draw(img)  
    imgR.rectangle(shape,  outline ="red",width=5) ##bb for animal

    confidence = bbxyxy[4]
    string_bb = 'animal ' + str(round(confidence, 2))
    font = ImageFont.truetype(FONTS[font_style], font_size) 

    text_size = font.getbbox(string_bb) # (h,w)
    position = (bbxyxy[0],bbxyxy[1] - text_size[1] -2 )
    left, top, right, bottom = imgR.textbbox(position, string_bb, font=font)
    imgR.rectangle((left, top-5, right+5, bottom+5), fill="red")
    imgR.text((bbxyxy[0] + 3 ,bbxyxy[1] - text_size[1] -2 ), string_bb, font=font, fill="black")

    return imgR

###########################################
def save_results_as_json(md_results, dlc_outputs, animal_bboxes, map_dlc_label_id_to_str, model,mega_model_input, path_to_output_file = 'download_predictions.json'):

    """
    Output detections as json file

    animal_bboxes: detection rows [x1,y1,x2,y2,conf,label], one per entry of dlc_outputs
    """
    # initialise dict to save to json
    info = {}
    info['date'] = str(today)
    info['MD_model'] = str(mega_model_input)
    # info from megaDetector
    info['file']= md_results.files[0]
    number_bb = len(md_results.xyxy[0].tolist())
    info['number_of_bb'] = number_bb
    # info from DLC
    info['dlc_model'] = model
    labels = [n for n in map_dlc_label_id_to_str.values()]

    # define aux dict for every animal bounding box above threshold
    for i in range(len(dlc_outputs)):
        aux={}
        # MD output
        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  animal_bboxes[i]
        aux['corner_1'] = (corner_x1,corner_y1)
        aux['corner_2'] = (corner_x2,corner_y2)
        aux['predict MD'] = md_results.names[0]
        aux['confidence MD'] = confidence
        
        # DLC output
        kypts = []
        for s in dlc_outputs[i]:
            aux1 = []
            for j in s:
                aux1.append(float(j))

            kypts.append(aux1)
        aux['dlc_pred']  = dict(zip(labels,kypts))
        info['bb_' + str(i) ]=aux

    # save dict as json
    with open(path_to_output_file, 'w') as f:
        json.dump(info, f, indent=1)
        print('Output file saved at {}'.format(path_to_output_file))

    return path_to_output_file


def save_results_only_dlc(dlc_outputs,map_label_id_to_str,model,output_file = 'dowload_predictions_dlc.json'):

    """
    write json dlc output
    """
    info = {}
    info['date'] = str(today)
    labels = [n for n in map_label_id_to_str.values()]
    info['dlc_model'] = model
    kypts = []
    for s in dlc_outputs:
        aux1 = []
        for j in s:
            aux1.append(float(j))

        kypts.append(aux1)
    info['dlc_pred']  = dict(zip(labels,kypts))

    with open(output_file, 'w') as f:
        json.dump(info, f, indent=1)
        print('Output file saved at {}'.format(output_file))

    return output_file


def save_results_pytorch(animals, map_label_id_to_str, model, pose_model, detector, path_to_output_file = 'download_predictions.json'):

    """
    Output PyTorch SuperAnimal predictions as json file (same layout as save_results_as_json)

    animals: list of {'bbox': [x1,y1,x2,y2,conf], 'kpts': (num_keypoints, 3)}, in image coords
    detector: None if the detector was skipped (whole image used as one animal)
    """
    info = {}
    info['date'] = str(today)
    info['backend'] = 'pytorch'
    info['dlc_model'] = model
    info['pose_model'] = pose_model
    info['detector'] = detector
    info['number_of_bb'] = len(animals)
    labels = [n for n in map_label_id_to_str.values()]

    for i, animal in enumerate(animals):
        corner_x1, corner_y1, corner_x2, corner_y2, confidence = animal['bbox']
        aux = {}
        aux['corner_1'] = (corner_x1, corner_y1)
        aux['corner_2'] = (corner_x2, corner_y2)
        aux['confidence'] = confidence
        aux['dlc_pred'] = dict(zip(labels, [[float(v) for v in kpt] for kpt in animal['kpts']]))
        info['bb_' + str(i)] = aux

    with open(path_to_output_file, 'w') as f:
        json.dump(info, f, indent=1)
        print('Output file saved at {}'.format(path_to_output_file))

    return path_to_output_file


###########################################