"""
Data ingestion logic
"""
import os
import random
import re
from io import BytesIO
from pathlib import Path
from typing import Optional, List, Dict

import numpy as np
import tensorflow as tf
from astropy.io import fits


def fits_open(path):
    return fits.open(str(path))[0].data.squeeze()


def transform(image, flip, seed, scale_size, crop_size):
    r = image
    if flip:
        r = tf.image.random_flip_left_right(r, seed=seed)
        r = tf.image.random_flip_up_down(r, seed=seed+1)

    # area produces a nice downscaling, but does nearest neighbor for upscaling
    # assume we're going to be doing downscaling here
    if scale_size and crop_size:
        if scale_size > crop_size:
            r = tf.image.resize_images(r, [scale_size, scale_size], method=tf.image.ResizeMethod.AREA)
            offset = tf.cast(tf.floor(tf.random_uniform([2], 0, scale_size - crop_size + 1, seed=seed)), dtype=tf.int32)
            r = tf.image.crop_to_bounding_box(r, offset[0], offset[1], crop_size, crop_size)
        elif scale_size < crop_size:
            raise Exception("scale size cannot be less than crop size")
    return r


def load_data(path,             # type: str
              crop_size=256,    # type: Optional[int]
              flip=False,       # type: bool
              scale_size=None,  # type: Optional[int]
              max_epochs=1,     # type: int
              batch_size=1,     # type: int
              loop=False,       # type: bool
              start=None,       # type: Optional[int]
              end=None          # type: Optional[int]
              ):                # type: (...) -> (tf.data.Dataset, int)
    """
    Point this to a path containing fits files fallowing naming schema <number>-<type>.fits

    Returns: a tensorflow dataset generator of dimensions [batch_size, x, y, channels]
    """

    p = Path(path)

    # find out our range
    r = re.compile('.*/(\d+)-skymodel.fits')
    ints = sorted([int(r.match(str(i)).group(1)) for i in p.glob("*-skymodel.fits")])
    if len(ints) > 1:
        min_, max_ = ints[0], ints[-1]
    else:
        min_, max_ = 0, 1

    if start:
        if min_ <= start < max_:
            min_ = start
        else:
            raise Exception("start ({}) out of range {} - {}".format(start, min_, max_))

    if end:
        if min_ < end <= max_ + 1:
            max_ = end
        else:
            raise Exception("end ({}) out of range {} - {}".format(end, min_, max_))

    count = max_ - min_

    def dataset_generator():
        for i in range(min_, max_):
            # todo: we need to encode the header as a string
            # header = fits.open("{}/{}-skymodel.fits".format(path, i))[0].header
            psf = fits_open("{}/{}-psf.fits".format(path, i))[:, :, np.newaxis]
            dirty = fits_open("{}/{}-dirty.fits".format(path, i))[:, :, np.newaxis]
            skymodel = fits_open("{}/{}-skymodel.fits".format(path, i))[:, :, np.newaxis]
            min_flux = dirty.min()
            max_flux = dirty.max()
            yield i, min_flux, max_flux, psf, dirty, skymodel

    ds = tf.data.Dataset.from_generator(dataset_generator,
                                        output_shapes=((), (), (), (512, 512, 1), (256, 256, 1), (256, 256, 1)),
                                        output_types=(tf.int32, tf.float32, tf.float32,
                                                      tf.float32, tf.float32, tf.float32),
                                        )

    # transforming
    # synchronize seed for image operations so that we do the same operations on all
    seed = random.randint(0, 2 ** 31 - 1)
    t = lambda i: transform(i, flip, seed, scale_size, crop_size)
    ds = ds.map(lambda i, mn, mx, psf, drty, skmd: (i, mn, mx, t(psf), t(drty), t(skmd)))

    if loop:
        ds = ds.repeat(max_epochs)

    ds = ds.batch(batch_size)
    return ds, count


def fits_decode(content):
    def internal(data):
        return fits.open(BytesIO(data))[0].data.squeeze().astype(np.float32)

    return tf.py_func(internal, [content], tf.float32)


def fits_encode(image):
    def internal(data):
        buf = BytesIO()
        hdu = fits.PrimaryHDU(data.squeeze())
        hdul = fits.HDUList([hdu])
        hdul.writeto(buf)
        return buf.getvalue()

    return tf.py_func(internal, [image], tf.string)


def save_images(fetches, output_dir, step=None, subfolder="images", extention="png"):
    # type: (Dict, str, Optional[int], Optional[str], Optional[str]) -> List[Dict]
    if subfolder:
        image_dir = os.path.join(output_dir, subfolder)
    else:
        image_dir = output_dir
    if not os.path.exists(image_dir):
        os.makedirs(image_dir)

    filesets = []
    indexs = fetches.pop('indexs')
    for key, value in fetches.items():
        for index, contents in zip(indexs, value):
            filename = str(index) + "-" + key + "." + extention
            filesets.append({"name": filename, "step": step})
            if step is not None:
                filename = "%08d-%s" % (step, filename)
            out_path = os.path.join(image_dir, filename)
            with open(out_path, 'wb') as f:
                f.write(contents)
    return filesets


def preprocess(image, min_, max_):
    with tf.name_scope("preprocess"):
        return (image / (max_/2.0)[:, None, None, None]) - 1
        #return (image / (max_ / 1.9)[:, None, None, None]) - 0.95


def deprocess(image, min_, max_):
    with tf.name_scope("deprocess"):
        return (image + 1) * (max_/2.0)[:, None, None, None]
        # return (image + 0.95) * (max_ / 1.9)[:, None, None, None]


