#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