#export
from pathlib import Path, PosixPath
import mimetypes
import pandas as pd
import os
import PIL.Image

from collections.abc import Iterable

from torch.utils.data import DataLoader

from functools import partial

import numpy as np
import torch

import matplotlib.pyplot as plt

path = PosixPath('/home/eross/.fastai/data/imagenette-160')

#export
Path.ls = lambda self: list(self.iterdir())

#export
def _get_files(p, fs, extensions=None):
    p = Path(p)
    res = [p/f for f in fs if not f.startswith('.')
           and ((not extensions) or f'.{f.split(".")[-1].lower()}' in extensions)]
    return res
                
def get_files(path, recurse=False, extensions=None, include=None):
    path = Path(path)
    extensions = set(extensions)
    extensions = {e.lower() for e in extensions}
    if recurse:
        res = []
        for i,(p,d,f) in enumerate(os.walk(path)): # returns (dirpath, dirnames, filenames)
            if include is not None and i==0: d[:] = [o for o in d if o in include]
            else:                            d[:] = [o for o in d if not o.startswith('.')]
            res += _get_files(p, f, extensions)
        return res
    else:
        f = [o.name for o in os.scandir(path) if o.is_file()]
        return _get_files(path, f, extensions)
                
def extensions_mime(ext):
    return frozenset(k for k,v in mimetypes.types_map.items() if v.startswith(ext))
                
EXTENSIONS_IMAGE = extensions_mime('image/')
EXTENSIONS_AUDIO = extensions_mime('audio/')
EXTENSIONS_VIDEO = extensions_mime('video/')
                
def img_get_files(path, recurse=False, extensions=None, include=None):
    if extensions is None: extensions = EXTENSIONS_IMAGE
    return get_files(path, recurse, extensions, include)

#export
# Compose and apply: Called compose in dl2v3 lectures
def comply(functions, x, reverse=False):
    if reverse: functions = reversed(functions)
    for f in functions:
        x = f(x)
    return x

#export
img_open = PIL.Image.open

def img_rgb(item): return item.convert('RGB')

def img_resize(item, size, method=PIL.Image.LANCZOS):
    return item.resize(size if isinstance(size, Iterable) else (size, size),
                       method)

def img_to_float(x): return torch.from_numpy(np.array(x, dtype=np.float32, copy=False)).permute(2,0,1).contiguous()/255.

def img_from_float(x, mode="RGB"): 
    c, h, w = x.shape
    assert c == len(mode), f"Mode {mode} doesn't agree with channels {c}"
    return PIL.Image.frombytes(mode, (w,h), bytes((x*255.).permute(1,2,0).contiguous().view(-1).byte()))

assert torch.allclose(img_to_float(img), img_to_float(img_from_float(img_to_float(img))))

assert torch.allclose(img_to_float(img), img_to_float(img_rgb(img)))

img_tfms = [img_open, partial(img_resize, size=128), img_rgb, img_to_float]

def _img_plot(ax, img, title=None):
    ax.axis('off')
    ax.imshow(img.permute(1,2,0))
    if title is not None:
        ax.set_title(title)

#export
def img_batch_plot(batch, titles=None, nrows=2, ncols=2):
    fig, axs = plt.subplots(nrows=nrows, ncols=ncols)
    fig.tight_layout()
    axs = [ax for row in axs for ax in row] # Flatten
    titles = titles or [None for _ in axs]
    for img, ax, title in zip(batch, axs, titles):
        _img_plot(ax, img, title)

#export
class ImgNormalize:
    def __init__(self, mean=None, sd=None, batch=True):
        self.shape = (1, -1, 1, 1) if batch else (-1, 1, 1)
        self.mean = mean
        self.sd = sd
        
    def get_stats(self, items):
        assert self.mean is None and self.sd is None, "Stats already calculated"
        n = 0
        mean = 0
        sd = 0
        for item in items:
            n += 1
            item = item.view(item.shape[0], -1)
            mean += item.mean(1)
            sd += item.std(1)
        self.mean = (mean / n).view(self.shape)
        self.sd = (sd / n).view(self.shape)
        return self
        
    def __call__(self, item):
        return (item - self.mean) / self.sd
    
    def invert(self, item):
        return item * self.sd + self.mean

assert torch.allclose(batch, normalize.invert(normalize(batch)), atol=1e-7)

#export
def is_unique(x):
    return len(x) == len(set(x))

#export
class CatNormalize:
    def __init__(self, c2i=None):
        self.c2i = c2i
        if c2i is not None and is_unique(c2i.values()):
            self.i2c = {i:c for c, i in c2i.items()}
    
    def get_stats(self, items):
        assert self.c2i is None, 'Already Initialised'
        self.c2i = {c:i for i, c in enumerate(set(items))}
        self.i2c = {i:c for c, i in self.c2i.items()}
        return self
        
    def __call__(self, item, batch=False):
        if not batch: item = [item]
        ans = [self.c2i[i] for i in item]
        if not batch: ans = ans[0]
        return torch.tensor(ans, dtype=torch.long)
    
    def invert(self, item, batch=False):
        if not batch: item = [item]
        ans = [self.i2c[i.item()] for i in item]
        if not batch: ans = ans[0]
        return ans

assert (cat_norm.invert(cat_norm(x, batch=True), batch=True) == x).all()

cat_tfms = [cat_norm]

cat_map = {
    0: 'chainsaw',
    1: 'golfball',
    2: 'petrol pump',
    3: 'trumbone',
    4: 'speaker',
    5: 'church',
    6: 'parachute',
    7: 'dog',
    8: 'tench',
    9: 'truck'
}

def batch_get_cat(items):
    return [cat_map[item.item()] for item in items]

# Kind of a flatmap for tensors
class TransformCollate:
    def __init__(self, tfms=[], collate=torch.stack):
        self.tfms = tfms
        self.collate = collate
            
    def __call__(self, items):
        return self.collate([comply(self.tfms, item) for item in items])

img_collate = TransformCollate(img_tfms)

bs=8
train_x = DataLoader(df.path[df.train], bs, collate_fn=img_collate, shuffle=True)

cat_collate = TransformCollate(cat_tfms)

train_y = DataLoader(df.label[df.train], bs, collate_fn=cat_collate,
                     shuffle=True # This will shuffle in a different way to the x's - which completely breaks it
                    )

class ProductCollate:
    def __init__(self, *collates):
        self.collates = collates
            
    def __call__(self, items):
        items = list(zip(*items))
        assert len(items) == len(self.collates)
        return tuple(collate(item_group) for collate, item_group in zip(self.collates, items))

xy_collate = ProductCollate(img_collate, cat_collate)

train = DataLoader(list(zip(df.path[df.train], df.label[df.train])),
                   bs,
                   collate_fn=xy_collate,
                   shuffle=True
                  )

def listify(p=None, q=None):
    "Make `p` listy and the same length as `q`."
    if p is None: p=[]
    elif isinstance(p, str):          p=[p]
    elif not isinstance(p, Iterable): p=[p]
    p = list(p)
    n = q if type(q)==int else len(p) if q is None else len(q)
    if len(p)==1: p = p * n
    assert len(p)==n, f'List len mismatch ({len(p)} vs {n})'
    return list(p)

class ListContainer():
    def __init__(self, items, tfms=None): self.items = listify(items); self.tfms = tfms or []
    def __getitem__(self, idx):
        if isinstance(idx, slice): return self.__class__(self.items[idx], self.tfms)
        try: return comply(self.tfms, self.items[idx])
        except TypeError:
            if isinstance(idx[0],(bool, np.bool_)):
                assert len(idx)==len(self) # bool mask
                return self.__class__([o for m,o in zip(idx,self.items) if m], self.tfms)
            return self.__class__([self.items[i] for i in idx], self.tfms)
        
    def get_inverse(self, idxs):
        if idxs is None: return self
        if isinstance(idxs, slice): idxs = range(len(self))[idxs]
        assert len(idxs) > 0, 'Must pass range of indices'
        if isinstance(idxs[0], (bool, np.bool_)):
            return self[[not x for x in idxs]]
        else:
            return self[[x for x in range(len(self)) if x not in idxs]]
        
    def __len__(self): return len(self.items)
    def __iter__(self): return (comply(self.tfms, x) for x in iter(self.items))
    def __setitem__(self, i, o): self.items[i] = o
    def __delitem__(self, i): del(self.items[i])
    # TODO: Check out repr lib in Python
    def __repr__(self):
        res = f'{self.__class__.__name__} ({len(self)} items)\n{self.items[:10]}'
        if len(self)>10: res = res[:-1]+ '...]'
        return res

def project(f, n):
    return lambda xs: [x if i!=n else f(x) for i, x in enumerate(xs)]

# Projection on Cartesian Product
def product_tfms(*tfmses):
    return [project(tfm, i) for i, tfms in enumerate(tfmses) for tfm in tfms]

xy_tfms = product_tfms(img_tfms, cat_tfms)

def project(f, n):
    return lambda xs: [x if i!=n else f(x) for i, x in enumerate(xs)]

# Projection on Cartesian Product
def product_tfms(*tfmses):
    return [project(tfm, i) for i, tfms in enumerate(tfmses) for tfm in tfms]

xy_tfms = product_tfms(img_tfms, cat_tfms)