import torch
import torchvision
from torchvision import transforms
import random
from PIL import Image
import os

from utils import RandomAffineAndRetMat

def load_filenames(data_dir):
  img_exts = ['.jpg', '.jpeg', '.png', '.bmp', '.ppm', '.pgm', '.tif', '.tiff']
  filenames = [f for f in os.listdir(data_dir) if os.path.splitext(f)[1].lower() in img_exts]

  return filenames


class MyDataset:
  def __init__(self, X, valid=False, img_dir='resources/trainB/', img_size=256):
    self.X = X
    self.valid = valid
    self.img_dir = img_dir
    self.img_size = img_size

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

  def __getitem__(self, index):
      # 画像を読み込んでトランスフォームを適用
    f = self.img_dir + self.X[index]
    original_X = Image.open(f)
    trans = [
      transforms.ToTensor(),
      # transforms.Normalize(mean=means, std=stds),
      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),

      transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.15),
      transforms.RandomGrayscale(0.3),
    ]
    transform = transforms.Compose(trans)
    xlist = []
    matlist = []
    is_flip = random.randint(0, 1) # 同じ画像はフリップ
    for i in range(2):
      af = RandomAffineAndRetMat(
        degrees=[-30, 30],
        translate=(0.1, 0.1), scale=(0.8, 1.2),
        # fill=(random.random(), random.random(), random.random()),
        fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
        shear=[-10, 10],
        interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
      )
      X, affine_matrix = af(transforms.Resize(self.img_size)(original_X))

      # randomflip
      if is_flip == 1:
        X = transforms.RandomHorizontalFlip(1.)(X)
        flip_matrix = torch.tensor([[-1., 0., 0.],
                                     [0., 1., 0.],
                                     [0., 0., 1.]])
        affine_matrix = torch.matmul(flip_matrix, affine_matrix)

      xlist.append(transform(X))
      matlist.append(affine_matrix)

    X = torch.stack(xlist)
    mat = torch.stack(matlist)
    return X, mat, f