# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use

import os, pdb
import numpy as np
from PIL import Image

from .dataset import Dataset
from .pair_dataset import PairDataset, StillPairDataset


class AachenImages (Dataset):
    """ Loads all images from the Aachen Day-Night dataset 
    """
    def __init__(self, select='db day night', root='data/aachen'):
        Dataset.__init__(self)
        self.root = root
        self.img_dir = 'images_upright'
        self.select = set(select.split())
        assert self.select, 'Nothing was selected'
        
        self.imgs = []
        root = os.path.join(root, self.img_dir)
        for dirpath, _, filenames in os.walk(root):
            r = dirpath[len(root)+1:]
            if not(self.select & set(r.split('/'))): continue
            self.imgs += [os.path.join(r,f) for f in filenames if f.endswith('.jpg')]
        
        self.nimg = len(self.imgs)
        assert self.nimg, 'Empty Aachen dataset'

    def get_key(self, idx):
        return self.imgs[idx]



class AachenImages_DB (AachenImages):
    """ Only database (db) images.
    """
    def __init__(self, **kw):
        AachenImages.__init__(self, select='db', **kw)
        self.db_image_idxs = {self.get_tag(i) : i for i,f in enumerate(self.imgs)}
    
    def get_tag(self, idx): 
        # returns image tag == img number (name)
        return os.path.split( self.imgs[idx][:-4] )[1]



class AachenPairs_StyleTransferDayNight (AachenImages_DB, StillPairDataset):
    """ synthetic day-night pairs of images 
        (night images obtained using autoamtic style transfer from web night images)
    """
    def __init__(self, root='data/aachen/style_transfer', **kw):
        StillPairDataset.__init__(self)
        AachenImages_DB.__init__(self, **kw)
        old_root = os.path.join(self.root, self.img_dir)
        self.root = os.path.commonprefix((old_root, root))
        self.img_dir = ''

        newpath = lambda folder, f: os.path.join(folder, f)[len(self.root):]
        self.imgs = [newpath(old_root, f) for f in self.imgs]

        self.image_pairs = []
        for fname in os.listdir(root):
            tag = fname.split('.jpg.st_')[0]
            self.image_pairs.append((self.db_image_idxs[tag], len(self.imgs)))
            self.imgs.append(newpath(root, fname))

        self.nimg = len(self.imgs)
        self.npairs = len(self.image_pairs)
        assert self.nimg and self.npairs



class AachenPairs_OpticalFlow (AachenImages_DB, PairDataset):
    """ Image pairs from Aachen db with optical flow.
    """
    def __init__(self, root='data/aachen/optical_flow', **kw):
        PairDataset.__init__(self)
        AachenImages_DB.__init__(self, **kw)
        self.root_flow = root

        # find out the subsest of valid pairs from the list of flow files
        flows = {f for f in os.listdir(os.path.join(root, 'flow')) if f.endswith('.png')}
        masks = {f for f in os.listdir(os.path.join(root, 'mask')) if f.endswith('.png')}
        assert flows == masks, 'Missing flow or mask pairs'
        
        make_pair = lambda f: tuple(self.db_image_idxs[v] for v in f[:-4].split('_'))
        self.image_pairs = [make_pair(f) for f in flows]
        self.npairs = len(self.image_pairs)
        assert self.nimg and self.npairs

    def get_mask_filename(self, pair_idx):
        tag_a, tag_b = map(self.get_tag, self.image_pairs[pair_idx])
        return os.path.join(self.root_flow, 'mask', f'{tag_a}_{tag_b}.png')

    def get_mask(self, pair_idx):
        return np.asarray(Image.open(self.get_mask_filename(pair_idx)))

    def get_flow_filename(self, pair_idx):
        tag_a, tag_b = map(self.get_tag, self.image_pairs[pair_idx])
        return os.path.join(self.root_flow, 'flow', f'{tag_a}_{tag_b}.png')

    def get_flow(self, pair_idx):
        fname = self.get_flow_filename(pair_idx)
        try:
            return self._png2flow(fname)
        except IOError:
            flow = open(fname[:-4], 'rb')
            help = np.fromfile(flow, np.float32, 1)
            assert help == 202021.25
            W, H = np.fromfile(flow, np.int32, 2)
            flow = np.fromfile(flow, np.float32).reshape((H, W, 2))
            return self._flow2png(flow, fname)

    def get_pair(self, idx, output=()):
        if isinstance(output, str): 
            output = output.split()

        img1, img2 = map(self.get_image, self.image_pairs[idx])
        meta = {}
        
        if 'flow' in output or 'aflow' in output:
            flow = self.get_flow(idx)
            assert flow.shape[:2] == img1.size[::-1]
            meta['flow'] = flow
            H, W = flow.shape[:2]
            meta['aflow'] = flow + np.mgrid[:H,:W][::-1].transpose(1,2,0)
        
        if 'mask' in output:
            mask = self.get_mask(idx)
            assert mask.shape[:2] == img1.size[::-1]
            meta['mask'] = mask
        
        return img1, img2, meta




if __name__ == '__main__':
    print(aachen_db_images)
    print(aachen_style_transfer_pairs)
    print(aachen_flow_pairs)
    pdb.set_trace()