import torch
import torchvision
import numpy as np
from torchtools.transforms import SmartCrop
import math

class Bucketeer():
    def __init__(self, dataloader, density=256*256, factor=8, ratios=[1/1, 1/2, 3/4, 3/5, 4/5, 6/9, 9/16], reverse_list=True, randomize_p=0.3, randomize_q=0.2, crop_mode='random', p_random_ratio=0.0, interpolate_nearest=False):
        assert crop_mode in ['center', 'random', 'smart']
        self.crop_mode = crop_mode
        self.ratios = ratios
        if reverse_list:
            for r in list(ratios):
                if 1/r not in self.ratios:
                    self.ratios.append(1/r)
        self.sizes = {}
        for dd in density:
          self.sizes[dd]= [(int(((dd/r)**0.5//factor)*factor), int(((dd*r)**0.5//factor)*factor)) for r in ratios]   
    
        self.batch_size = dataloader.batch_size
        self.iterator = iter(dataloader)
        all_sizes =  []
        for k, vs in self.sizes.items():
          all_sizes += vs
        self.buckets = {s: [] for s in all_sizes}
        self.smartcrop = SmartCrop(int(density**0.5), randomize_p, randomize_q) if self.crop_mode=='smart' else None
        self.p_random_ratio = p_random_ratio
        self.interpolate_nearest = interpolate_nearest

    def get_available_batch(self):
        for b in self.buckets:
            if len(self.buckets[b]) >= self.batch_size:
                batch = self.buckets[b][:self.batch_size]
                self.buckets[b] = self.buckets[b][self.batch_size:]
                return batch
        return None

    def get_closest_size(self, x):
        w, h = x.size(-1), x.size(-2)

            
        best_size_idx = np.argmin([abs(w/h-r) for r in self.ratios])
        find_dict = {dd : abs(w*h - self.sizes[dd][best_size_idx][0]*self.sizes[dd][best_size_idx][1]) for dd, vv in self.sizes.items()}
        min_ = find_dict[list(find_dict.keys())[0]]
        find_size = self.sizes[list(find_dict.keys())[0]][best_size_idx]
        for dd, val in find_dict.items():
          if val < min_:
            min_ = val
            find_size = self.sizes[dd][best_size_idx]  
            
        return find_size

    def get_resize_size(self, orig_size, tgt_size):
        if (tgt_size[1]/tgt_size[0] - 1) * (orig_size[1]/orig_size[0] - 1) >= 0:
            alt_min = int(math.ceil(max(tgt_size)*min(orig_size)/max(orig_size)))
            resize_size = max(alt_min, min(tgt_size))
        else:
            alt_max = int(math.ceil(min(tgt_size)*max(orig_size)/min(orig_size)))
            resize_size = max(alt_max, max(tgt_size))

        return resize_size

    def __next__(self):
        batch = self.get_available_batch()
        while batch is None:
            elements = next(self.iterator)
            for dct in elements:
                img = dct['images']
                size = self.get_closest_size(img)
                resize_size = self.get_resize_size(img.shape[-2:], size)
              
                if self.interpolate_nearest:
                    img = torchvision.transforms.functional.resize(img, resize_size, interpolation=torchvision.transforms.InterpolationMode.NEAREST)
                else:
                    img = torchvision.transforms.functional.resize(img, resize_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR, antialias=True)
                if self.crop_mode == 'center':
                    img = torchvision.transforms.functional.center_crop(img, size)
                elif self.crop_mode == 'random':
                    img = torchvision.transforms.RandomCrop(size)(img)
                elif self.crop_mode == 'smart':
                    self.smartcrop.output_size = size
                    img = self.smartcrop(img)
                
                self.buckets[size].append({**{'images': img}, **{k:dct[k] for k in dct if k != 'images'}})
            batch = self.get_available_batch()

        out = {k:[batch[i][k] for i in range(len(batch))] for k in batch[0]}
        return {k: torch.stack(o, dim=0) if isinstance(o[0], torch.Tensor) else o for k, o in out.items()}