"""
The io module contains various utility methods for loading/saving Python objects/dicts in various formats.
Currently, there is support for the most frequent data storing formats:
    - JSON (+gzip)
    - pickled Python objects (+gzip)
"""

import gzip
import io
import json
import pickle
from io import BytesIO
from pathlib import Path
from typing import Union, Tuple, Optional

from matplotlib import pyplot as plt

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal

import PIL.Image
import cv2
import imageio
import numpy as np
import requests
import yaml
from PIL import Image

from elias.util.fs import ensure_file_ending, ensure_directory_exists_for_file

PathType = Union[str, Path]


# =========================================================================
# Zipped JSON (.json.gz)
# =========================================================================

def save_zipped_json(obj: dict, path: PathType, suffix: str = 'json.gz'):
    """
    Parses the given Python dict into json, zips it and stores it at the specified location.
    Per default, the path will have a suffix 'json.gz'.

    Parameters
    ----------
        obj: dict
            The JSON (represented as python dict) to be zipped and stored
        path: str
            Where to store the file
        suffix: str, default 'json.gz'
            File name suffix
    """

    path = ensure_file_ending(path, suffix)
    ensure_directory_exists_for_file(path)
    with gzip.open(path, 'wb') as f:
        json.dump(obj, f)


def load_zipped_json(path: PathType, suffix: str = 'json.gz') -> dict:
    """
    Loads the specified compressed JSON file from `path` and extracts it.
    Per default, the path name is assumed to have a suffix 'json.gz'.

    Parameters
    ----------
        path: str
            Path to the zipped JSON to load
        suffix: str, default 'json.gz'
            File name suffix

    Returns
    -------
        The contents of the zipped JSON as a Python dict
    """

    path = ensure_file_ending(path, suffix)
    with gzip.open(path, 'rb') as f:
        return json.load(f)


# =========================================================================
# Zipped pickles (.p.gz)
# =========================================================================

def save_zipped_object(obj: object, path: PathType, suffix: str = 'p.gz'):
    """
    Pickles, zips and stores an arbitrary python object at the specified `path`.
    Per default, the file will have a suffix 'json.gz'.

    Parameters
    ----------
        obj: object
            The Python object to be zipped and stored
        path:
            Where to store the file
        suffix: str, default 'p.gz'
            File name suffix
    """

    path = ensure_file_ending(path, suffix)
    ensure_directory_exists_for_file(path)
    with gzip.open(path, 'wb') as f:
        pickle.dump(obj, f)


def load_zipped_object(path: PathType, suffix: str = 'p.gz') -> object:
    """
    Loads zipped and pickled Python object stored at `path`.
    Per default, the file name is assumed to have a suffix 'p.gz'.

    Parameters
    ----------
        path: str
            Path to the zipped Python object that should be loaded
        suffix: str, default 'p.gz'
            File name suffix

    Returns
    -------
        The un-pickled Python object
    """

    path = ensure_file_ending(path, suffix)
    with gzip.open(path, 'rb') as f:
        return pickle.load(f)


# =========================================================================
# Pickled objects (.p)
# =========================================================================

def save_pickled(obj: object, path: PathType, suffix: str = 'p'):
    """
    Pickles and stores an arbitrary python object at the specified `path`.
    Per default, the file will have a suffix 'p'.

    Parameters
    ----------
        obj: object
            The Python object to be pickled and stored
        path: str
            Where to store the file
        suffix: str, default 'p'
            File name suffix
    """

    path = ensure_file_ending(path, suffix)
    ensure_directory_exists_for_file(path)
    with open(f"{path}", 'wb') as f:
        pickle.dump(obj, f)


def load_pickled(path: PathType, suffix: str = 'p') -> object:
    """
    Loads a pickled Python object stored at the specified `path`.
    Per default, the file name is assumed to have a suffix 'p'.

    Parameters
    ----------
        path: str
            Path to the pickled Python object that should be loaded
        suffix: str, default 'p'
            File name suffix

    Returns
    -------
        The un-pickled Python object
    """

    path = ensure_file_ending(path, suffix)
    with open(path, 'rb') as f:
        return pickle.load(f)


# =========================================================================
# JSON (.json)
# =========================================================================

class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


def save_json(obj: dict, path: PathType, suffix: str = 'json'):
    """
    Stores the given Python dict in JSON format at `path`.
    Per default, the file will have a suffix 'json'.

    Parameters
    ----------
        obj: dict
            The Python dict to store
        path: str
            Where to store the file
        suffix:
            File name suffix
    """

    path = ensure_file_ending(path, suffix)
    ensure_directory_exists_for_file(path)
    with open(path, 'w') as f:
        json.dump(obj, f, indent=4, cls=NumpyEncoder)


def load_json(path: PathType, suffix: str = 'json') -> dict:
    """
    Loads and parses the given JSON file and returns it as a Python dict.
    Per default, the file name is assumed to have a suffix 'json'.

    Parameters
    ----------
        path: str
            Path to the JSON file
        suffix: str, default 'json'
            File name suffix

    Returns
    -------
        Contents of the JSON file as Python dict
    """

    path = ensure_file_ending(path, suffix)
    with open(path, 'r') as f:
        return json.load(f)


# =========================================================================
# YAML (.yaml)
# =========================================================================

def save_yaml(obj: dict, path: PathType, suffix: str = 'yaml'):
    path = ensure_file_ending(path, suffix)
    ensure_directory_exists_for_file(path)
    with open(path, 'w') as f:
        yaml.dump(obj, f)


def load_yaml(path: PathType, suffix: str = 'yaml') -> dict:
    path_with_suffix = ensure_file_ending(path, suffix)
    if not Path(path_with_suffix).exists():
        if suffix == 'yaml':
            # Try loading .yml instead
            return load_yaml(path, 'yml')
        else:
            raise FileNotFoundError(f"Could not load YAML file {path_with_suffix} (Neither with .yaml suffix). "
                                    f"Is the path correct?")

    with open(path_with_suffix, 'r') as f:
        return yaml.load(f, Loader=yaml.FullLoader)


# =========================================================================
# Images
# =========================================================================

def save_img(img: np.ndarray, path: PathType):
    """
    Save the given numpy array as an 8-bit image. The image type is determined by the file extension.
    `image` should be a np.uint8 array with shape [H, W, (C)] with C in {1, 3, 4} and values in range [0, 255]
    If `image` has type np.float32, it is assumed that values are in range [0, 1]. The float image will be automatically converted to uint8 type.

    Parameters
    ----------
        img: image as a [H, W, (C)] numpy array with values in [0, 255] or [0, 1].
        path: where to store the image. Any necessarily folders will be created automatically.
    """

    ensure_directory_exists_for_file(path)

    if Path(path).suffix == '.exr':
        imageio.imwrite(path, img)
    else:
        if img.dtype == np.float32:
            # TODO: Do we really want to quantize without asking?
            # If float array was passed, have to transform [0.0, 1.0] -> [0, ..., 255)
            assert 0 <= img.min() and img.max() <= 1, \
                "passed float array should have values between 0 and 1 to be interpreted as image"
            img = (img * 255).astype(np.uint8)

        if len(img.shape) == 3 and img.shape[2] == 1:
            # Assume this should be stored as a single-channel grayscale image
            img = img[:, :, 0]

        Image.fromarray(img).save(path)


def load_img(path: PathType) -> np.ndarray:
    if Path(path).suffix == '.exr':
        img = imageio.imread_v2(path)
    else:
        img = Image.open(path)

    return np.asarray(img)


def download_img(url: str) -> np.ndarray:
    response = requests.get(url)
    pil_img = Image.open(BytesIO(response.content))
    img = np.asarray(pil_img)

    return img


InterpolationType = Literal['nearest', 'bilinear', 'bicubic', 'lanczos']


def resize_img(img: np.ndarray,
               scale: Union[float, Tuple[float, float]],
               interpolation: InterpolationType = 'bilinear',
               use_opencv: bool = False) -> np.ndarray:
    try:
        iter(scale)
    except TypeError:
        scale_x = scale
        scale_y = scale
    else:
        scale_x, scale_y = scale

    if use_opencv:
        if interpolation == 'nearest':
            interpolation = cv2.INTER_NEAREST
        elif interpolation == 'bilinear':
            interpolation = cv2.INTER_LINEAR
        elif interpolation == 'bicubic':
            interpolation = cv2.INTER_CUBIC
        elif interpolation == 'lanczos':
            interpolation = cv2.INTER_LANCZOS4
        else:
            raise ValueError(f"Invalid interpolation type: {interpolation}")

        img = cv2.resize(img, (int(img.shape[1] * scale_x), int(img.shape[0] * scale_y)), interpolation=interpolation)
    else:
        if interpolation == 'nearest':
            interpolation = PIL.Image.Resampling.NEAREST
        elif interpolation == 'bilinear':
            interpolation = PIL.Image.Resampling.BILINEAR
        elif interpolation == 'bicubic':
            interpolation = PIL.Image.Resampling.BICUBIC
        elif interpolation == 'lanczos':
            interpolation = PIL.Image.Resampling.LANCZOS
        else:
            raise ValueError(f"Invalid interpolation type: {interpolation}")

        img_pil = Image.fromarray(img)
        img_pil = img_pil.resize((int(img.shape[1] * scale_x), int(img.shape[0] * scale_y)), resample=interpolation)
        img = np.array(img_pil)

    return img


def fig2img(fig: Optional[plt.Figure] = None, transparent: bool = False) -> np.ndarray:
    """
    Converts a Matplotlib figure to a numpy image.
    If no figure is given, the current open matplot figure will be used instead.
    """

    buf = io.BytesIO()
    if fig is None:
        fig = plt.gcf()

    fig.savefig(buf, bbox_inches='tight', transparent=transparent)
    buf.seek(0)
    img = np.asarray(Image.open(buf))

    if transparent:
        return img[..., :4]
    else:
        return img[..., :3]
