# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01e_utils.io.ipynb.

# %% auto 0
__all__ = ['image_extensions', 'save_txt', 'open_txt', 'open_json', 'save_json', 'mkdir', 'extract_files', 'get_files',
           'get_image_files', 'open_image', 'open_mask', 'save_image', 'RenderDict']

# %% ../../nbs/01e_utils.io.ipynb 2
import shutil, tempfile
import numpy as np
import polvo as pv
import uuid, json, mimetypes
from pathlib import Path
from fastcore.all import *
from PIL import Image, ImageOps

# %% ../../nbs/01e_utils.io.ipynb 3
def save_txt(s, path, append=False):
    mode = 'a' if append else 'w'
    with open(str(path), mode) as f: f.write(s)

# %% ../../nbs/01e_utils.io.ipynb 4
def open_txt(path):
    with open(str(path), 'r') as f: return f.read()
#     return np.loadtxt(path, dtype=str, delimiter="\n").tolist()

# %% ../../nbs/01e_utils.io.ipynb 6
def open_json(path):
    with open(str(path)) as f: return json.load(f)

# %% ../../nbs/01e_utils.io.ipynb 7
def save_json(d, path):
    with open(str(path), 'w') as f: json.dump(d, f)

# %% ../../nbs/01e_utils.io.ipynb 9
def mkdir(path, exist_ok=False, parents=False, overwrite=False, tmp=False, **kwargs):
    "Creates and returns a directory defined by `path`, optionally removing previous existing directory if `overwrite` is `True`"
    path = Path(path)
    if tmp:
        path = tempfile.gettempdir() / path
        parents = True
    if path.exists() and overwrite: shutil.rmtree(path)
    path.mkdir(exist_ok=exist_ok, parents=parents, **kwargs)
    return path

# %% ../../nbs/01e_utils.io.ipynb 11
def extract_files(files, extract_to_dir, show_pbar=True):
    for file in pbar(files, show=show_pbar):
        extract_path = extract_to_dir / Path(file).with_suffix("").name
        shutil.unpack_archive(file, extract_path)
    return extract_to_dir

# %% ../../nbs/01e_utils.io.ipynb 12
def _get_files(p, fs, extensions=None):
    "COPIED FROM https://github.com/fastai/fastai/blob/master/nbs/05_data.transforms.ipynb"
    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

# %% ../../nbs/01e_utils.io.ipynb 13
def get_files(path, extensions=None, recurse=True, folders=None, followlinks=True):
    """Get all the files in `path` with optional `extensions`, optionally with `recurse`, only in `folders`, if specified.
    COPIED FROM https://github.com/fastai/fastai/blob/master/nbs/05_data.transforms.ipynb
    """
    path = Path(path)
    folders=L(folders)
    extensions = setify(extensions)
    extensions = {e.lower() for e in extensions}
    if recurse:
        res = []
        for i,(p,d,f) in enumerate(os.walk(path, followlinks=followlinks)): # returns (dirpath, dirnames, filenames)
            if len(folders) !=0 and i==0: d[:] = [o for o in d if o in folders]
            else:                         d[:] = [o for o in d if not o.startswith('.')]
            if len(folders) !=0 and i==0 and '.' not in folders: continue
            res += _get_files(p, f, extensions)
    else:
        f = [o.name for o in os.scandir(path) if o.is_file()]
        res = _get_files(path, f, extensions)
    return L(res)

# %% ../../nbs/01e_utils.io.ipynb 15
image_extensions = set(k for k,v in mimetypes.types_map.items() if v.startswith('image/'))

# %% ../../nbs/01e_utils.io.ipynb 16
def get_image_files(path, recurse=True, folders=None):
    """Get image files in `path` recursively, only in `folders`, if specified.
    COPIED FROM https://github.com/fastai/fastai/blob/master/nbs/05_data.transforms.ipynb
    """
    return get_files(path, extensions=image_extensions, recurse=recurse, folders=folders)

# %% ../../nbs/01e_utils.io.ipynb 18
def open_image(path, gray=False, ignore_exif=True) -> Image.Image:
    "Open an image from disk `path` as a PIL Image"
    color = "L" if gray else "RGB"
    image = Image.open(str(path))
    if not ignore_exif: image = ImageOps.exif_transpose(image)
    return image.convert(color)

# %% ../../nbs/01e_utils.io.ipynb 20
@delegates(open_image)
def open_mask(path, gray=True, **kwargs) -> Image.Image:
    return open_image(path, gray=gray, **kwargs)

# %% ../../nbs/01e_utils.io.ipynb 22
def save_image(image, path):
    if isinstance(image, np.ndarray): image = Image.fromarray(image)
    return image.save(str(path))

# %% ../../nbs/01e_utils.io.ipynb 24
class RenderDict:
    "From https://www.reddit.com/r/IPython/comments/34t4m7/lpt_print_json_in_collapsible_format_in_ipython/"
    def __init__(self, json_data):
        if isinstance(json_data, dict):
            self.json_str = json.dumps(json_data)
        else:
            self.json_str = json_data
        self.uuid = str(uuid.uuid4())

    def _ipython_display_(self):
        from IPython.display import display_javascript, display_html, display
        display_html('<div id="{}" style="height: 600px; width:100%;"></div>'.format(self.uuid), raw=True)
        display_javascript("""
        require(["https://rawgit.com/caldwell/renderjson/master/renderjson.js"], function() {
        document.getElementById('%s').appendChild(renderjson(%s))
        });
        """ % (self.uuid, self.json_str), raw=True)
