import os
from pathlib import Path
from PIL import Image

import torch
import torchvision
from torchvision.ops import MultiScaleRoIAlign
from torchvision.datasets import FakeData
from torchvision.transforms.functional import pil_to_tensor
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.backbone_utils import _resnet_fpn_extractor
from torchvision.models.detection.faster_rcnn import (
    FasterRCNN_ResNet50_FPN_Weights, FastRCNNPredictor, FasterRCNN as _FasterRCNN
)
from torchvision.models.detection import FasterRCNN
from torchvision.models.resnet import resnet50, ResNet50_Weights
from torch import nn


class TestingDataset(torch.utils.data.Dataset):
    def __init__(self, img_root):
        """
        Create a FakeData dataset for testing
        Args:
            img_root: path to the image directory.
        """
        img_root.mkdir(exist_ok=True)
        self.img_root = img_root

        self.img_names = sorted(os.listdir(img_root))

        # generate a fake dataset of images if it doesn't already exist.
        # for testing inference, we don't need labels/bboxes, just images
        if len(self.img_names) == 0:
            fake_dset = FakeData(size=1000, image_size=(3,256,256))
            for idx, (img, _) in enumerate(fake_dset):
                img.save(Path(img_root / f"img{idx}.png"), "PNG")
            self.img_names = sorted(os.listdir(img_root))

    def __len__(self):
        return len(self.img_names)

    def __getitem__(self, idx):
        img_name = self.img_names[idx]
        img_path = os.path.join(self.img_root, img_name)
        img = Image.open(img_path).convert("RGB")

        img = pil_to_tensor(img).to(dtype=torch.float32)
        img /= 255.
        return (img,)# label)

def get_fasterrcnn(num_classes):
    # Define new anchor sizes and aspect ratios
    # You can customize these values based on your requirements
    new_anchor_sizes = ((4,), (8,), (16,), (32,), (64,),)
    new_aspect_ratios = (
        (0.5, 1.0, 2.0, 4.0, 8.0),
    ) * len(new_anchor_sizes)

    # Create a new AnchorGenerator with custom anchor sizes and aspect ratios
    anchor_generator = AnchorGenerator(
        sizes=new_anchor_sizes,
        aspect_ratios=new_aspect_ratios,
    )
    print('Anchors', anchor_generator.num_anchors_per_location())

    roi_pooler = torchvision.ops.MultiScaleRoIAlign(
        featmap_names=['0'],
        output_size=7,
        sampling_ratio=2,
    )

    resnet = resnet50(
        weights=ResNet50_Weights.IMAGENET1K_V1, 
        norm_layer=nn.BatchNorm2d
    )
    resnet_fpn = _resnet_fpn_extractor(resnet, 5)

    model = _FasterRCNN(
        resnet_fpn,
        num_classes=num_classes,
        min_size=256,
        max_size=256,
        rpn_anchor_generator=anchor_generator,
        box_roi_pool=roi_pooler,
    )

    return model

data_root = Path("./data/random_test_imgs")
data_root.mkdir(exist_ok=True, parents=True)

dataset = TestingDataset(
    data_root / 'images'
)
dataloader = torch.utils.data.DataLoader(
    dataset,
    # threshold is batchsize 172: larger batches throw bus error or seg
    # fault, smaller output bounding boxes, labels, scores as expected
    batch_size=172,
    shuffle=False,
    # num_workers=4,
)
it = iter(dataloader)

num_classes = 2
model = get_fasterrcnn(num_classes)
model.eval()
with torch.no_grad():
    # import pdb; pdb.set_trace()
    # foo = model(next(it)[0].to(device))
    predictions = model(next(it)[0])
print(
    predictions[0]['boxes'][:5],
    predictions[0]['labels'][:5],
    predictions[0]['scores'][:5]
)