import os

import cv2
import numpy as np
import torch
from skimage import io
from skimage.transform import resize
from torch.utils.data import Dataset

from saicinpainting.evaluation.evaluator import InpaintingEvaluator
from saicinpainting.evaluation.losses.base_loss import SSIMScore, LPIPSScore, FIDScore


class SimpleImageDataset(Dataset):
    def __init__(self, root_dir, image_size=(400, 600)):
        self.root_dir = root_dir
        self.files = sorted(os.listdir(root_dir))
        self.image_size = image_size

    def __getitem__(self, index):
        img_name = os.path.join(self.root_dir, self.files[index])
        image = io.imread(img_name)
        image = resize(image, self.image_size, anti_aliasing=True)
        image = torch.FloatTensor(image).permute(2, 0, 1)
        return image

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


def create_rectangle_mask(height, width):
    mask = np.ones((height, width))
    up_left_corner = width // 4, height // 4
    down_right_corner = (width - up_left_corner[0] - 1, height - up_left_corner[1] - 1)
    cv2.rectangle(mask, up_left_corner, down_right_corner, (0, 0, 0), thickness=cv2.FILLED)
    return mask


class Model():
    def __call__(self, img_batch, mask_batch):
        mean = (img_batch * mask_batch[:, None, :, :]).sum(dim=(2, 3)) / mask_batch.sum(dim=(1, 2))[:, None]
        inpainted = mean[:, :, None, None] * (1 - mask_batch[:, None, :, :]) + img_batch * mask_batch[:, None, :, :]
        return inpainted


class SimpleImageSquareMaskDataset(Dataset):
    def __init__(self, dataset):
        self.dataset = dataset
        self.mask = torch.FloatTensor(create_rectangle_mask(*self.dataset.image_size))
        self.model = Model()

    def __getitem__(self, index):
        img = self.dataset[index]
        mask = self.mask.clone()
        inpainted = self.model(img[None, ...], mask[None, ...])
        return dict(image=img, mask=mask, inpainted=inpainted)

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


dataset = SimpleImageDataset('imgs')
mask_dataset = SimpleImageSquareMaskDataset(dataset)
model = Model()
metrics = {
    'ssim': SSIMScore(),
    'lpips': LPIPSScore(),
    'fid': FIDScore()
}

evaluator = InpaintingEvaluator(
    mask_dataset, scores=metrics, batch_size=3, area_grouping=True
)

results = evaluator.evaluate(model)
print(results)