import json
from collections import defaultdict

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from torchmetrics.functional.classification import binary_average_precision
from tqdm import tqdm

from constants import *
from DenseAV.denseav.shared import unnorm, remove_axes


def prep_heatmap(sims, masks, h, w):
    masks = masks.to(torch.float32)
    hm = torch.einsum("bhwt,bt->bhw", sims, masks) / masks.sum(-1).reshape(-1, 1, 1)
    hm -= hm.min()
    hm /= hm.max()
    return F.interpolate(hm.unsqueeze(1), (h, w), mode="bilinear").squeeze(1)


def iou(prediction, target):
    prediction = prediction > 0.0
    target = target > 0.5
    intersection = torch.logical_and(prediction, target).sum().float()
    union = torch.logical_or(prediction, target).sum().float()
    if union == 0:
        return 1.0
    return (intersection / union).item()  # Convert to Python scalar


def multi_iou(prediction, target, k=20):
    prediction = torch.tensor(prediction)
    target = torch.tensor(target)
    target = target > 0.5

    thresholds = torch.linspace(prediction.min(), prediction.max(), k)
    hard_pred = prediction.unsqueeze(0) > thresholds.reshape(k, 1, 1, 1, 1)
    target = torch.broadcast_to(target.unsqueeze(0), hard_pred.shape)

    # Calculate IoU for each threshold
    intersection = torch.logical_and(hard_pred, target).sum(dim=(1, 2, 3, 4)).float()
    union = torch.logical_or(hard_pred, target).sum(dim=(1, 2, 3, 4)).float()
    union = torch.where(union == 0, torch.tensor(1.0), union)  # Avoid division by zero
    iou_scores = intersection / union

    # Find the best IoU and corresponding threshold
    best_iou, best_idx = torch.max(iou_scores, dim=0)
    # best_threshold = thresholds[best_idx]
    # print(best_threshold)
    return best_iou  # , best_threshold.item()


def get_paired_heatmaps(
        model,
        results,
        class_ids,
        timing,
        class_names=None):
    sims = model.sim_agg.get_pairwise_sims(
        results,
        raw=False,
        agg_sim=False,
        agg_heads=True
    ).squeeze(1).mean(-2)

    prompt_classes = torch.tensor(list(class_ids))
    gt = results["semseg"] == prompt_classes.reshape(-1, 1, 1)
    basic_masks = results[AUDIO_MASK]  # BxT
    _, fullh, fullw = gt.shape
    basic_heatmaps = prep_heatmap(sims, basic_masks, fullh, fullw)

    if timing is not None:
        prompt_timing = np.array(list(timing))
        raw_timing = torch.tensor([json.loads(t) for t in prompt_timing])
        timing = torch.clone(raw_timing)
        timing[:, 0] -= .2
        timing[:, 1] += .2
        total_length = (results['total_length'] / 16000)[0]
        fracs = timing / total_length
        bounds = basic_masks.shape[1] * fracs
        bounds[:, 0] = bounds[:, 0].floor()
        bounds[:, 1] = bounds[:, 1].ceil()
        bounds = bounds.to(torch.int64)
        advanced_masks = (F.one_hot(bounds, basic_masks.shape[1]).cumsum(-1).sum(-2) == 1).to(basic_masks)
        advanced_heatmaps = prep_heatmap(sims, advanced_masks, fullh, fullw)

    metrics = defaultdict(list)
    unique_classes = torch.unique(prompt_classes)

    should_plot = class_names is not None

    if should_plot:
        prompt_names = np.array(list(class_names))

    for prompt_class in tqdm(unique_classes):
        subset = torch.where(prompt_classes == prompt_class)[0]
        gt_subset = gt[subset]
        basic_subset = basic_heatmaps[subset]
        metrics["basic_ap"].append(binary_average_precision(basic_subset.flatten(), gt_subset.flatten()))
        metrics["basic_iou"].append(multi_iou(basic_subset.flatten(), gt_subset.flatten()))

        if timing is not None:
            advanced_subset = advanced_heatmaps[subset]
            metrics["advanced_ap"].append(binary_average_precision(advanced_subset.flatten(), gt_subset.flatten()))
            metrics["advanced_iou"].append(multi_iou(advanced_subset.flatten(), gt_subset.flatten()))

        if should_plot:
            prompt_class_subset = prompt_classes[subset]
            name_subset = prompt_names[subset]
            print(prompt_class, name_subset, prompt_class_subset)
            n_imgs = min(len(subset), 5)
            if n_imgs > 1:
                fig, axes = plt.subplots(n_imgs, 5, figsize=(4 * 5, n_imgs * 3))
                frame_subset = unnorm(results[IMAGE_INPUT][subset].squeeze(1)).permute(0, 2, 3, 1)
                semseg_subset = results["semseg"][subset]
                for img_num in range(n_imgs):
                    axes[img_num, 0].imshow(frame_subset[img_num])
                    axes[img_num, 1].imshow(basic_subset[img_num])
                    axes[img_num, 2].imshow(advanced_subset[img_num])
                    axes[img_num, 3].imshow(gt_subset[img_num])
                    axes[img_num, 4].imshow(semseg_subset[img_num], cmap="tab20", interpolation='none')

                axes[0, 0].set_title("Image")
                class_name = name_subset[0].split(",")[0]
                axes[0, 1].set_title(f"{class_name} Basic Heatmap")
                axes[0, 2].set_title(f"{class_name} Advanced Heatmap")
                axes[0, 3].set_title("True Mask")
                axes[0, 4].set_title("True Seg")
                remove_axes(axes)
                plt.tight_layout()
                plt.show()

    return metrics, unique_classes