import datetime
import numpy as np
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import collections as mc
import matplotlib.pyplot as plt
import math
from aris import BEAM_WIDTH_DIR
import cv2

from dataloader import create_dataloader_aris


STANDARD_FIG_SIZE = (16, 9)
OUT_PDF_FILE_NAME = 'multipage_pdf.pdf'


def make_pdf(i, state, result, table_headers):

    fish_info = result["fish_info"][i]
    fish_table = result["fish_table"][i]
    json_result = result['json_result'][i]
    metadata = json_result['metadata']
    aris_input = result["aris_input"][i]

    with PdfPages(OUT_PDF_FILE_NAME) as pdf:
        plt.rcParams['text.usetex'] = False
        
        generate_title_page(pdf, metadata, state)
        
        generate_global_result(pdf, fish_info)

        generate_fish_list(pdf, table_headers, fish_table)


        dataset = None
        if (aris_input is not None):
            dataloader, dataset = create_dataloader_aris(aris_input, BEAM_WIDTH_DIR, None)

        for i, fish in enumerate(json_result['fish']):
            calculate_fish_paths(json_result, dataset, i)
        
        draw_combined_fish_graphs(pdf, json_result)

        for i, fish in enumerate(json_result['fish']):
            generate_fish_tracks(pdf, json_result, i)

        # We can also set the file's metadata via the PdfPages object:
        d = pdf.infodict()
        d['Title'] = 'Multipage PDF Example'
        d['Author'] = 'Oskar Åström'
        d['Subject'] = 'How to create a multipage pdf file and set its metadata'
        d['Keywords'] = ''
        d['CreationDate'] = datetime.datetime.today()
        d['ModDate'] = datetime.datetime.today()


def generate_title_page(pdf, metadata, state):
    # set up figure that will be used to display the opening banner
    fig = plt.figure(figsize=STANDARD_FIG_SIZE)
    plt.axis('off')

    title_font_size = 40
    minor_font_size = 20

    # stuff to be printed out on the first page of the report
    plt.text(0.5,-0.5,f'{metadata["FILE_NAME"].split("/")[-1]}',fontsize=title_font_size, horizontalalignment='center')

    plt.text(0,1,f'Duration: {metadata["TOTAL_TIME"]}',fontsize=minor_font_size)
    plt.text(0,1.5,f'Frames: {metadata["TOTAL_FRAMES"]}',fontsize=minor_font_size)
    plt.text(0,2,f'Frame Rate: {metadata["FRAME_RATE"]}',fontsize=minor_font_size)

    plt.text(0.5,1,f'Time of filming: {metadata["DATE"]} ({metadata["START"]} - {metadata["END"]})',fontsize=minor_font_size)
    plt.text(0.5,1.5,f'Web app version: {state["version"]}',fontsize=minor_font_size)

    plt.text(1.1,4.5,f'PDF generated on {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',fontsize=minor_font_size, horizontalalignment='right')

    plt.ylim([-1, 4])
    plt.xlim([0, 1])
    plt.gca().invert_yaxis()
    
    pdf.savefig(fig)
    plt.close(fig)

def generate_global_result(pdf, fish_info):
    # set up figure that will be used to display the opening banner
    fig = plt.figure(figsize=STANDARD_FIG_SIZE)
    plt.axis('off')
    # stuff to be printed out on the first page of the report

    minor_font_size = 18

    headers = ["Result", "Camera Info", "Hyperparameters"]
    info_col_1 = []
    info_col_2 = []
    info_col = info_col_1
    row_state = -1
    for row in fish_info:
        if row_state >= 0:
            info_col.append([row[0].replace("**","").replace("_", " ").lower(), row[1], 'normal'])
        if (row[0] == "****"):
            row_state += 1
            if row_state == 2: info_col = info_col_2
            info_col.append([headers[row_state], "", 'bold'])
    for row_i, row in enumerate(info_col_1):
        h = -1 + 5*row_i/len(info_col_1)
        plt.text(0,    h, row[0], fontsize=minor_font_size, weight=row[2])
        plt.text(0.25, h, row[1], fontsize=minor_font_size, weight=row[2])
    for row_i, row in enumerate(info_col_2):
        h = -1 + 5*row_i/len(info_col_2)
        plt.text(0.5,  h, row[0], fontsize=minor_font_size, weight=row[2])
        plt.text(0.75, h, row[1], fontsize=minor_font_size, weight=row[2])
    plt.ylim([-1, 4])
    plt.xlim([0, 1])
    plt.gca().invert_yaxis()
    
    pdf.savefig(fig)
    plt.close(fig)

def generate_fish_list(pdf, table_headers, fish_table):
    # set up figure that will be used to display the opening banner
    fig = plt.figure(figsize=STANDARD_FIG_SIZE)
    plt.axis('off')
    # stuff to be printed out on the first page of the report

    title_font_size = 40
    header_font_size = 12
    body_font_size = 20

    # Title
    plt.text(0.5,-1.3,f'{"Identified Fish"}',fontsize=title_font_size, horizontalalignment='center', weight='bold')

    # Identified fish
    row_h = 0.25
    col_start = 0
    row_l = 1
    dropout_i = None
    for col_i, col in enumerate(table_headers):
        x = col_start + row_l*(col_i+0.5)/len(table_headers)
        if col == "TOTAL": col = "ID"
        if col == "DETECTION_DROPOUT": 
            col = "frame drop rate"
            dropout_i = col_i
        col = col.lower().replace("_", " ")
        plt.text(x, -1, col, fontsize=header_font_size, horizontalalignment='center', weight="bold")
        plt.plot([col_start*2, -col_start*2 + row_l], [-1 + 0.05, -1 + 0.05], color='black')

    for row_i, row in enumerate(fish_table):
        y = -1 + (row_i+1)*row_h
        plt.plot([col_start*2, -col_start*2 + row_l], [y+0.05, y+0.05], color='black')
        for col_i, col in enumerate(row):
            x = col_start + row_l*(col_i+0.5)/len(row)
            if (col_i == dropout_i and type(col) is not str):
                col = str(int(col*100)) + "%"
            elif type(col) == float:
                col = "{:.4f}".format(col)
            plt.text(x, y, col, fontsize=body_font_size, horizontalalignment='center')
    plt.ylim([-1, 4])
    plt.xlim([0, 1])
    plt.gca().invert_yaxis()
    
    pdf.savefig(fig)
    plt.close(fig)

def calculate_fish_paths(result, dataset, id):

    fish = result['metadata']['FISH'][id]
    start_frame = fish['START_FRAME']
    end_frame = fish['END_FRAME']

    # Extract base frame (first frame for that fish)
    w, h = 1, 2
    img = None
    if (dataset is not None):
        
        images = dataset.didson.load_frames(start_frame=start_frame, end_frame=start_frame+1)
        img = images[0]

        frame_height = 2
        scale_factor = frame_height / h
        h = frame_height
        w = int(scale_factor*w)

    fish['base_frame'] = img
    fish['scaled_frame_size'] = (w, h)


    # Find frames for this fish
    bboxes = []
    for frame in result['frames'][start_frame:end_frame+1]:
        bbox = None
        for ann in frame['fish']:
            if ann['fish_id'] == id+1:
                bbox = ann
        bboxes.append(bbox)


    # Calculate tracks through frames
    missed = 0
    X = []
    Y = []
    V = []
    certainty = []
    for bbox in bboxes:
        if bbox is not None:

            # Find fish centers
            x = (bbox['bbox'][0] + bbox['bbox'][2])/2
            y = (bbox['bbox'][1] + bbox['bbox'][3])/2

            # Calculate velocity
            v = None
            if len(X) > 0:
                last_x = X[-1]
                last_y = Y[-1]
                dx = result['image_meter_width']*(last_x - x)/(missed+1)
                dy = result['image_meter_height']*(last_y - y)/(missed+1)
                v = math.sqrt(dx*dx + dy*dy)

            # Interpolate over missing frames
            if missed > 0:    
                for i in range(missed):
                    p = (i+1)/(missed+1)
                    X.append(last_x*(1-p) + p*x)
                    Y.append(last_y*(1-p) + p*y)
                    V.append(v)
                    certainty.append(False)
            
            # Append new track frame
            X.append(x)
            Y.append(y)
            if v is not None: V.append(v)
            certainty.append(True)
            missed = 0
        else:
            missed += 1

    fish['path'] = {
        'X': X,
        'Y': Y,
        'certainty': certainty,
        'V': V
    }


def draw_combined_fish_graphs(pdf, result):
    vel = []
    log_vel = []
    for fish in result['metadata']['FISH']:
        vel += fish['path']['V']
        log_vel += [math.log(v) for v in fish['path']['V']]

    fig, axs = plt.subplots(2, 2, sharey=True, figsize=STANDARD_FIG_SIZE)
    axs[0,0].hist(log_vel, bins=20)
    axs[0,0].set_title('Fish Log-Velocities between frames')
    axs[0,0].set_xlabel("Log-Velocity (log(m/frame))")
    axs[0,1].hist(vel, bins=20)
    axs[0,1].set_title('Fish Velocities between frames')
    axs[0,1].set_xlabel("Velocity (m/frame)")

    pdf.savefig(fig)
    plt.close(fig)
    

def generate_fish_tracks(pdf, result, id):

    fish = result['metadata']['FISH'][id]
    start_frame = fish['START_FRAME']
    end_frame = fish['END_FRAME']

    fig, ax = plt.subplots(figsize=STANDARD_FIG_SIZE)
    plt.axis('off')

    w, h = fish['scaled_frame_size']
    if (fish['base_frame'] is not None):
        img = fish['base_frame']
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
        plt.imshow(img, extent=(0, h, 0, w), cmap=plt.colormaps['Greys'])

    # Title
    plt.text(h/2,1.1,f'Fish {id+1} (frames {start_frame}-{end_frame})',fontsize=40, color="red", horizontalalignment='center', zorder=5)

    X = fish['path']['X']
    Y = fish['path']['Y']
    certainty = fish['path']['certainty']
    
    plt.text(h*(1-Y[0]), w*(1-X[0]), "Start", fontsize=15, weight="bold")
    plt.text(h*(1-Y[-1]), w*(1-X[-1]), "End", fontsize=15, weight="bold")

    colors = [""]
    for i in range(1, len(X)):

        certain = certainty[i]
        fully_certain = certain
        half_certain = certain
        if i > 0: 
            fully_certain &= certainty[i-1]
            half_certain |= certainty[i-1]
            
        #color = 'yellow' if certain else 'orangered'
        #plt.plot(h*(1-y), w*(1-x), marker='o', markersize=3, color=color, zorder=3)
        col = 'yellow' if fully_certain else ('darkorange' if half_certain else 'orangered')
        colors.append(col)
        ax.plot([h*(1-Y[i-1]), h*(1-Y[i])], [w*(1-X[i-1]), w*(1-X[i])], color=col, linewidth=1)
    
    print(len(X))
    print(len(Y))
    print(len(colors))
    for i in range(1, len(X)):
        ax.plot(h*(1-Y[i]), w*(1-X[i]), color=colors[i], marker='o', markersize=3)
    

    plt.ylim([0, w])
    plt.xlim([0, h])
    pdf.savefig(fig)
    plt.close(fig)