import variables
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import os
import random
import torch

from collections import Counter
from torch.utils.data import DataLoader
from tqdm import tqdm


def iou_width_height(boxes1, boxes2):
    """
    Parameters:
        boxes1 (tensor): width and height of the first bounding boxes
        boxes2 (tensor): width and height of the second bounding boxes
    Returns:
        tensor: Intersection over union of the corresponding boxes
    """
    intersection = torch.min(boxes1[..., 0], boxes2[..., 0]) * torch.min(
        boxes1[..., 1], boxes2[..., 1]
    )
    union = (
        boxes1[..., 0] * boxes1[..., 1] + boxes2[..., 0] * boxes2[..., 1] - intersection
    )
    return intersection / union


def IOU(frames_pred, frames_labels, frame_format="1_point"):

    if frame_format == "1_point":
        box1_x1 = frames_pred[..., 0:1] - frames_pred[..., 2:3] / 2
        box1_y1 = frames_pred[..., 1:2] - frames_pred[..., 3:4] / 2
        box1_x2 = frames_pred[..., 0:1] + frames_pred[..., 2:3] / 2
        box1_y2 = frames_pred[..., 1:2] + frames_pred[..., 3:4] / 2
        box2_x1 = frames_labels[..., 0:1] - frames_labels[..., 2:3] / 2
        box2_y1 = frames_labels[..., 1:2] - frames_labels[..., 3:4] / 2
        box2_x2 = frames_labels[..., 0:1] + frames_labels[..., 2:3] / 2
        box2_y2 = frames_labels[..., 1:2] + frames_labels[..., 3:4] / 2

    if frame_format == "2_points":
        box1_x1 = frames_pred[..., 0:1]
        box1_y1 = frames_pred[..., 1:2]
        box1_x2 = frames_pred[..., 2:3]
        box1_y2 = frames_pred[..., 3:4]
        box2_x1 = frames_labels[..., 0:1]
        box2_y1 = frames_labels[..., 1:2]
        box2_x2 = frames_labels[..., 2:3]
        box2_y2 = frames_labels[..., 3:4]

    x1 = torch.max(box1_x1, box2_x1)
    y1 = torch.max(box1_y1, box2_y1)
    x2 = torch.min(box1_x2, box2_x2)
    y2 = torch.min(box1_y2, box2_y2)

    intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)
    box1_area = abs((box1_x2 - box1_x1) * (box1_y2 - box1_y1))
    box2_area = abs((box2_x2 - box2_x1) * (box2_y2 - box2_y1))

    return intersection / (box1_area + box2_area - intersection + 1e-6)


def NMS(frames, iou_lim, prob_lim, frame_format="1_point"):
    assert type(frames) == list

    frames = [box for box in frames if box[1] > prob_lim]
    frames = sorted(frames, key=lambda x: x[1], reverse=True)
    frames_after_nms = []

    while frames:
        chosen_box = frames.pop(0)

        frames = [
            frame
            for frame in frames
            if frame[0] != chosen_box[0]
               or IOU(
                torch.tensor(chosen_box[2:]),
                torch.tensor(frame[2:]),
                frame_format=frame_format,
            )
               < iou_lim
        ]

        frames_after_nms.append(chosen_box)

    return frames_after_nms


def MAP(pred_frames, true_frames, iou_lim=0.5, frame_format="1_point", no_classes=9):

    average_precisions = []

    eps = 1e-6

    for cls in range(no_classes):
        detections = []
        ground_truths = []

        for det in pred_frames:
            if det[1] == cls:
                detections.append(det)

        for true_box in true_frames:
            if true_box[1] == cls:
                ground_truths.append(true_box)

        no_frames = Counter([gt[0] for gt in ground_truths])

        for key, val in no_frames.items():
            no_frames[key] = torch.zeros(val)

        detections.sort(key=lambda x: x[2], reverse=True)
        True_positives = torch.zeros((len(detections)))
        False_positives = torch.zeros((len(detections)))
        total_no_frames = len(ground_truths)

        # If none exists for this class then we can safely skip
        if total_no_frames == 0:
            continue

        for index, det in enumerate(detections):
            # Only take out the ground_truths that have the same
            # training idx as det
            ground_truth_img = [
                bbox for bbox in ground_truths if bbox[0] == det[0]
            ]

            num_gts = len(ground_truth_img)
            top_iou = 0

            for idx, gt in enumerate(ground_truth_img):
                iou = IOU(
                    torch.tensor(det[3:]),
                    torch.tensor(gt[3:]),
                    frame_format=frame_format,
                )

                if iou > top_iou:
                    top_iou = iou
                    top_gt_idx = idx

            if top_iou > iou_lim:
                # only detect ground truth det once
                if no_frames[det[0]][top_gt_idx] == 0:
                    # true positive and add this bounding box to seen
                    True_positives[index] = 1
                    no_frames[det[0]][top_gt_idx] = 1
                else:
                    False_positives[index] = 1

            else:
                False_positives[index] = 1

        true_pos_cs = torch.cumsum(True_positives, dim=0)
        false_pos_cs = torch.cumsum(False_positives, dim=0)
        recalls = true_pos_cs / (total_no_frames + eps)
        recalls = torch.cat((torch.tensor([0]), recalls))
        precisions = true_pos_cs / (true_pos_cs + false_pos_cs + eps)
        precisions = torch.cat((torch.tensor([1]), precisions))
        average_precisions.append(torch.trapz(precisions, recalls))

    return sum(average_precisions) / len(average_precisions)


# TODO: Delete this function
def plot_image(image, boxes):

    cmap = plt.get_cmap("tab20b")
    class_labels = variables.CLASSES
    colors = [cmap(i) for i in np.linspace(0, 1, len(class_labels))]
    im = np.array(image)
    height, width, _ = im.shape

    # Create figure and axes
    fig, ax = plt.subplots(1)
    # Display the image
    ax.imshow(im)

    # box[0] is x midpoint, box[2] is width
    # box[1] is y midpoint, box[3] is height

    # Create a Rectangle patch
    for box in boxes:
        assert len(box) == 6, "box should contain class pred, confidence, x, y, width, height"
        class_pred = box[0]
        box = box[2:]
        upper_left_x = box[0] - box[2] / 2
        upper_left_y = box[1] - box[3] / 2
        rect = patches.Rectangle(
            (upper_left_x * width, upper_left_y * height),
            box[2] * width,
            box[3] * height,
            linewidth=2,
            edgecolor=colors[int(class_pred)],
            facecolor="none",
        )
        # Add the patch to the Axes
        ax.add_patch(rect)
        plt.text(
            upper_left_x * width,
            upper_left_y * height,
            s=class_labels[int(class_pred)],
            color="white",
            verticalalignment="top",
            bbox={"color": colors[int(class_pred)], "pad": 0},
        )

    plt.show()


def get_eval_frames(model, loader, iou_lim, prob_lim, anchors, frame_format="1_point"):
    model.eval()
    train_idx = 0
    all_pred_boxes = []
    all_true_boxes = []
    for batch_idx, (x, labels) in enumerate(tqdm(loader)):
        x = x.to(variables.DEVICE)

        with torch.no_grad():
            predictions = model(x)

        batch_size = x.shape[0]
        frames = [[] for _ in range(batch_size)]
        for i in range(3):
            SPLIT_SIZE = predictions[i].shape[2]
            anchor = torch.tensor([*anchors[i]]).to(variables.DEVICE) * SPLIT_SIZE
            boxes_scale_i = CTF(
                predictions[i], anchor, SPLIT_SIZE=SPLIT_SIZE, is_preds=True
            )
            for index, (frame) in enumerate(boxes_scale_i):
                frames[index] += frame

        # we just want one bbox for each label, not one for each scale
        true_bboxes = CTF(
            labels[2], anchor, SPLIT_SIZE=SPLIT_SIZE, is_preds=False
        )

        for index in range(batch_size):
            nms_boxes = NMS(
                frames[index],
                iou_lim=iou_lim,
                prob_lim=prob_lim,
                frame_format=frame_format,
            )

            for nms_box in nms_boxes:
                all_pred_boxes.append([train_idx] + nms_box)

            for box in true_bboxes[index]:
                if box[1] > prob_lim:
                    all_true_boxes.append([train_idx] + box)

            train_idx += 1

    model.train()
    return all_pred_boxes, all_true_boxes

# Cells to frames
def CTF(preds, anchors, SPLIT_SIZE, is_preds=True):

    BATCH_SIZE = preds.shape[0]
    num_anchors = len(anchors)
    frame_preds = preds[..., 1:5]
    if is_preds:
        anchors = anchors.reshape(1, len(anchors), 1, 1, 2)
        frame_preds[..., 0:2] = torch.sigmoid(frame_preds[..., 0:2])
        frame_preds[..., 2:] = torch.exp(frame_preds[..., 2:]) * anchors
        scores = torch.sigmoid(preds[..., 0:1])
        best_class = torch.argmax(preds[..., 5:], dim=-1).unsqueeze(-1)
    else:
        scores = preds[..., 0:1]
        best_class = preds[..., 5:6]

    cell_dets = (
        torch.arange(SPLIT_SIZE)
        .repeat(preds.shape[0], 3, SPLIT_SIZE, 1)
        .unsqueeze(-1)
        .to(preds.device)
    )
    x = 1 / SPLIT_SIZE * (frame_preds[..., 0:1] + cell_dets)
    y = 1 / SPLIT_SIZE * (frame_preds[..., 1:2] + cell_dets.permute(0, 1, 3, 2, 4))
    w_h = 1 / SPLIT_SIZE * frame_preds[..., 2:4]
    converted_frames = torch.cat((best_class, scores, x, y, w_h), dim=-1).reshape(BATCH_SIZE, num_anchors * SPLIT_SIZE * SPLIT_SIZE, 6)
    return converted_frames.tolist()

def check_class_accuracy(model, loader, threshold):
    model.eval()
    amount_class_preds, true_class = 0, 0
    amount_noobj, true_noobj = 0, 0
    amount_obj, true_obj = 0, 0

    for index, (x, y) in enumerate(tqdm(loader)):
        x = x.to(variables.DEVICE)
        with torch.no_grad():
            out = model(x)

        for i in range(3):
            y[i] = y[i].to(variables.DEVICE)
            obj = y[i][..., 0] == 1
            noobj = y[i][..., 0] == 0

            true_class += torch.sum(
                torch.argmax(out[i][..., 5:][obj], dim=-1) == y[i][..., 5][obj]
            )
            amount_class_preds += torch.sum(obj)

            cls_preds = torch.sigmoid(out[i][..., 0]) > threshold
            true_obj += torch.sum(cls_preds[obj] == y[i][..., 0][obj])
            amount_obj += torch.sum(obj)
            true_noobj += torch.sum(cls_preds[noobj] == y[i][..., 0][noobj])
            amount_noobj += torch.sum(noobj)

    model.train()


def get_mean_std(loader):
    # var[X] = E[X**2] - E[X]**2
    sum_channel, channels_sqrt_sum, no_batches = 0, 0, 0

    for data, _ in tqdm(loader):
        sum_channel += torch.mean(data, dim=[0, 2, 3])
        channels_sqrt_sum += torch.mean(data ** 2, dim=[0, 2, 3])
        no_batches += 1

    mean = sum_channel / no_batches
    std = (channels_sqrt_sum / no_batches - mean ** 2) ** 0.5

    return mean, std


def save_model(model, optimizer, filename="model.pth.tar"):
    try:
        print("--> Attempting to save a checkpoint")
        checkpoint = {
            "state_dict": model.state_dict(),
            "optimizer": optimizer.state_dict(),
        }
        torch.save(checkpoint, filename)
    except Exception as e:
        print(f"Saving checkpoint failed with message: {e}")


def load_model(model_file, model, optimizer, lr):
    print("=> Loading checkpoint")
    mod = torch.load(model_file, map_location=variables.DEVICE)
    model.load_state_dict(mod["state_dict"])
    optimizer.load_state_dict(mod["optimizer"])

    for parameters in optimizer.param_groups:
        parameters["lr"] = lr


def get_loaders(train_path, test_path):
    from dataset import KittiDataset

    IMG_SHAPE = variables.IMG_SHAPE
    train_dataset = KittiDataset(
        train_path,
        transform=variables.train_transforms,
        split_size=[IMG_SHAPE // 32, IMG_SHAPE // 16, IMG_SHAPE // 8],
        image_dir=variables.IMG_DIR,
        label_dir=variables.LABEL_DIR,
        anchors=variables.ANCHORS,
    )
    test_dataset = KittiDataset(
        test_path,
        transform=variables.test_transforms,
        split_size=[IMG_SHAPE // 32, IMG_SHAPE // 16, IMG_SHAPE // 8],
        image_dir=variables.IMG_DIR,
        label_dir=variables.LABEL_DIR,
        anchors=variables.ANCHORS,
    )
    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=variables.BATCH_SIZE,
        num_workers=variables.NO_WORKERS,
        pin_memory=variables.MEM,
        shuffle=True,
        drop_last=False,
    )
    test_loader = DataLoader(
        dataset=test_dataset,
        batch_size=variables.BATCH_SIZE,
        num_workers=variables.NO_WORKERS,
        pin_memory=variables.MEM,
        shuffle=False,
        drop_last=False,
    )

    train_eval_dataset = KittiDataset(
        train_path,
        transform=variables.test_transforms,
        split_size=[IMG_SHAPE // 32, IMG_SHAPE // 16, IMG_SHAPE // 8],
        image_dir=variables.IMG_DIR,
        label_dir=variables.LABEL_DIR,
        anchors=variables.ANCHORS,
    )
    train_eval_loader = DataLoader(
        dataset=train_eval_dataset,
        batch_size=variables.BATCH_SIZE,
        num_workers=variables.NO_WORKERS,
        pin_memory=variables.MEM,
        shuffle=False,
        drop_last=False,
    )

    return train_loader, test_loader, train_eval_loader



