import os
import re
from typing import Sequence, Tuple
import numpy as np
import torch
from torch.utils.data import DataLoader
from ..a3.annotation import read_groundtruth_file, AnnotationRect
from .label_grid import get_label_grid, draw_annotation_rects
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from .anchor_grid import get_anchor_grid
from PIL import Image
from torchvision.transforms import transforms
from itertools import islice


class MMP_Dataset(torch.utils.data.Dataset):
    def __init__(
        self,
        path_to_data: str,
        image_size: int,
        anchor_grid: np.ndarray,
        min_iou: float,
        is_test: bool,
    ):
        """
        @param anchor_grid: The anchor grid to be used for every image
        @param min_iou: The minimum IoU that is required for an overlap for the label grid.
        @param is_test: Whether this is the test set (True) or the validation/training set (False)
        """
        self.image_size = image_size
        self.images = []
        self.anchor_grid = anchor_grid
        self.min_iou = min_iou
        self.is_test = is_test
        img_pattern = re.compile(r"^(\d+)\.jpg$")
        files = set(os.listdir(path_to_data))

        for fname in files:
            match = img_pattern.match(fname)
            if match:
                img_file = os.path.join(path_to_data, fname)
                annotations = read_groundtruth_file(
                    os.path.join(path_to_data, f"{match.group(1)}.gt_data.txt")
                )
                self.images.append((img_file, annotations))

        self.images.sort(
            key=lambda x: int(re.match(r"(.*/)(\d+)(\.jpg)", x[0]).group(2))
        )

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, int]:
        """
        @return: 3-tuple of image tensor, label grid, and image (file-)number
        """
        img = Image.open(self.images[idx][0]).convert("RGB")
        padding = self.__padding__(img)
        transform = transforms.Compose(
            [
                transforms.Pad(padding, 0),
                transforms.Resize((self.image_size, self.image_size)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
            ]
        )
        img_tensor = transform(img)
        label_grid = get_label_grid(
            anchor_grid=self.anchor_grid, gts=self.images[idx][1], min_iou=self.min_iou
        )
        img_id = re.match(r".*(\/)([0-9]+)(\.[^\/]*$)", self.images[idx][0]).group(2)
        return (img_tensor, label_grid, int(img_id))

    def __len__(self) -> int:
        return len(self.images)

    def __padding__(self, img) -> Tuple[int, int, int, int]:
        w, h = img.size
        size = max(w, h)
        right_pad = size - w
        bottom_pad = size - h
        return (0, 0, right_pad, bottom_pad)


def get_dataloader(
    path_to_data: str,
    image_size: int,
    batch_size: int,
    num_workers: int,
    anchor_grid: np.ndarray,
    is_test: bool,
    is_train: bool = False,
) -> DataLoader:
    dataset = MMP_Dataset(
        path_to_data=path_to_data,
        image_size=image_size,
        is_test=is_test,
        anchor_grid=anchor_grid,
        min_iou=0.7,
    )
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=is_train,
        num_workers=num_workers,
        pin_memory=True,
    )
    return dataloader


def calculate_max_coverage(loader: DataLoader, min_iou: float) -> float:
    """
    @param loader: A DataLoader object, generated with the get_dataloader function.
    @param min_iou: Minimum IoU overlap that is required to count a ground truth box as covered.
    @return: Ratio of how mamy ground truth boxes are covered by a label grid box. Must be a value between 0 and 1.
    """
    raise NotImplementedError()


def print_img_tensor_with_annotations(
    img: torch.Tensor, annotations: Sequence["AnnotationRect"], output_file: str
):
    # Convert tensor to numpy, permute dimensions
    img_np = img.permute(1, 2, 0).cpu().numpy()
    img_np = img_np.astype(np.uint8)

    fig, ax = plt.subplots(1)
    ax.imshow(img_np)
    for rect in annotations:
        x1, y1, x2, y2 = rect.x1, rect.y1, rect.x2, rect.y2
        w = x2 - x1
        h = y2 - y1
        patch = patches.Rectangle(
            (x1, y1), w, h, linewidth=2, edgecolor="red", facecolor="none"
        )
        ax.add_patch(patch)
    plt.axis("off")
    plt.tight_layout(pad=0)
    plt.savefig(output_file, bbox_inches="tight", pad_inches=0)
    plt.close(fig)


def print_positive_boxes(
    img_tensor: torch.Tensor,
    label_grid: np.ndarray,
    img_id: torch.Tensor,
    anchor_grid: np.ndarray,
    path_to_data: str,
):
    annotations = [
        AnnotationRect.fromarray(anchor_grid[idx])
        for idx in np.ndindex(anchor_grid.shape[:-1])
        if label_grid[idx]
    ]
    print_img_tensor_with_annotations(
        img_tensor,
        annotations=annotations,
        output_file=f"mmp/a4/{img_id}_transformed.png",
    )
    draw_annotation_rects(
        annotations=annotations,
        image=f"{os.path.join(path_to_data, f'{str(img_id.item()).zfill(8)}.jpg')}",
        output_path=f"mmp/a4/{img_id}_original.png",
    )


def main():
    anchor_grid = get_anchor_grid(
        anchor_widths=[16, 32, 64, 96, 128, 144, 150, 160, 192, 224, 256],
        aspect_ratios=[1 / 3, 1 / 2, 3 / 5, 2 / 3, 3 / 4, 1, 4 / 3, 5 / 3, 2, 2.5, 3],
        num_rows=32,
        num_cols=32,
        scale_factor=20,
    )
    dataloader = get_dataloader(
        num_workers=6,
        is_train=True,
        is_test=False,
        batch_size=8,
        image_size=224,
        path_to_data=".data/mmp-public-3.2/train",
        anchor_grid=anchor_grid,
    )

    for img, label, img_id in islice(dataloader, 12):
        print_positive_boxes(
            img_tensor=img[5],
            label_grid=label[5],
            img_id=img_id[5],
            anchor_grid=anchor_grid,
            path_to_data=".data/mmp-public-3.2/train",
        )


if __name__ == "__main__":
    main()