# Copyright (C) 2015, Carlo de Franchis <carlo.de-franchis@cmla.ens-cachan.fr>
# Copyright (C) 2015, Gabriele Facciolo <facciolo@cmla.ens-cachan.fr>
# Copyright (C) 2015, Enric Meinhardt <enric.meinhardt@cmla.ens-cachan.fr>
# Copyright (C) 2015, Julien Michel <julien.michel@cnes.fr>


from __future__ import print_function
import os
import sys
import errno
import datetime
import warnings
import tempfile
import subprocess
import numpy as np
import rasterio


from s2p.config import cfg


# silent rasterio NotGeoreferencedWarning
warnings.filterwarnings("ignore",
                        category=rasterio.errors.NotGeoreferencedWarning)

# add the bin folder to system path
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
bin_dir = os.path.join(parent_dir, 'bin')
os.environ['PATH'] = bin_dir + os.pathsep + os.environ['PATH']

# global variable
# list of intermediary files generated by s2p
garbage = list()


def remove(target):
    try:
        os.remove(target)
    except OSError:
        pass

def garbage_cleanup():
    """
    Removes all the files listed in the global variable 'garbage'.
    """
    if cfg['clean_tmp']:
        while garbage:
            remove(garbage.pop())


def tmpfile(ext=''):
    """
    Creates a temporary file in the cfg['temporary_dir'] directory.

    Args:
        ext: desired file extension. The dot has to be included.

    Returns:
        absolute path to the created file

    The path of the created file is added to the garbage list to allow cleaning
    at the end of the pipeline.
    """
    fd, out = tempfile.mkstemp(suffix=ext, prefix='s2p_',
                               dir=os.path.expandvars(cfg['temporary_dir']))
    os.close(fd)           # http://www.logilab.org/blogentry/17873
    garbage.append(out)
    return out


class RunFailure(Exception):
    pass


def run(cmd, env=os.environ):
    """
    Runs a shell command, and print it before running.

    Arguments:
        cmd: string to be passed to a shell
        env (optional, default value is os.environ): dictionary containing the
            environment variables

    Both stdout and stderr of the shell in which the command is run are those
    of the parent process.
    """
    print("\nRUN: %s" % cmd)
    t = datetime.datetime.now()
    try:
        subprocess.check_call(cmd, shell=True, stdout=sys.stdout,
                              stderr=sys.stderr, env=env)
        print(datetime.datetime.now() - t)

    except subprocess.CalledProcessError as e:
        # raise a custom exception because the CalledProcessError causes the
        # pool to crash
        raise RunFailure({"command": e.cmd, "environment": env, "output":
                          e.output})


def mkdir_p(path):
    """
    Create a directory without complaining if it already exists.
    """
    try:
        os.makedirs(path)
    except OSError as exc: # requires Python > 2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else: raise


def shellquote(s):
    return "'%s'" % s.replace("'", "'\\''")


def matrix_translation(x, y):
    t = np.eye(3)
    t[0, 2] = x
    t[1, 2] = y
    return t


def image_size_gdal(im):
    """
    Read the width, height and pixel dimension of an image using rasterio.

    Args:
        im: path to the input image file

    Returns:
        w, h, pd: a tuple of length 3
    """
    with rasterio.open(im, 'r') as f:
        y, x = f.shape
        pd = f.count
    return x, y, pd


def gdal_read_as_array_with_nans(im):
    """
    Read an image replacing gdal NoDataValue with np.nan

    Args:
        im: path to the input image file

    Returns:
        array: raster as numpy array
    """
    with rasterio.open(im, 'r') as src:
        array = src.read()
        nodata_values = src.nodatavals

    for band, nodata in zip(array, nodata_values):
        if nodata is not None:
            band[band == nodata] = np.nan

    return array.squeeze()


def image_zoom_out_morpho(im, f):
    """
    Image zoom out by morphological operation (median).

    Args:
        im: path to the input image
        f: zoom out factor. It has to be a positive integer

    Returns:
        path to the output image
    """
    if (f != np.floor(f)):
        print('image_zoom_out_morpho: zoom factor has to be integer')
        sys.exit()

    out = tmpfile('.tif')
    run('downsa e %d %s %s' % (f, im, out))
    return out


def rasterio_write(path, array, profile={}, tags={}):
    """
    Write a numpy array in a tiff or png file with rasterio.

    Args:
        path (str): path to the output tiff/png file
        array (numpy array): 2D or 3D array containing the image to write.
        profile (dict): rasterio profile (ie dictionary of metadata)
        tags (dict): dictionary with additional geotiff tags
    """
    # determine the driver based on the file extension
    extension = os.path.splitext(path)[1].lower()
    if extension in ['.tif', '.tiff']:
        driver = 'GTiff'
    elif extension in ['.png']:
        driver = 'png'
    else:
        raise NotImplementedError('format {} not supported'.format(extension))

    # read image size and number of bands
    array = np.atleast_3d(array)
    height, width, nbands = array.shape

    # define image metadata dict
    profile.update(driver=driver, count=nbands, width=width, height=height,
                   dtype=array.dtype)

    # write to file
    with rasterio.Env():
        with rasterio.open(path, 'w', **profile) as dst:
            dst.write(np.transpose(array, (2, 0, 1)))
            dst.update_tags(**tags)


def image_apply_homography(out, im, H, w, h):
    """
    Applies an homography to an image.

    Args:
        out: path to the output image file
        im: path to the input image file
        H: numpy array containing the 3x3 homography matrix
        w, h: dimensions (width and height) of the output image

    The output image is defined on the domain [0, w] x [0, h]. Its pixels
    intensities are defined by out(x) = im(H^{-1}(x)).

    This function calls the homography binary, rewritten by Marc Lebrun and
    Carlo de Franchis based on a code of Pascal Monasse refactored by Gabriele
    Facciolo.
    """
    # write the matrix to a string
    hij = " ".join(str(x) for x in H.flatten())

    # apply the homography
    run("homography %s -h \"%s\" %s %d %d" % (im, hij, out, w, h))


def image_qauto(im, out=None):
    """
    Uniform requantization between min and max intensity.

    Args:
        im: path to input image
        out (optional, default is None): path to output image

    Returns:
        path of requantized image, saved as png
    """
    if out is None:
        out = tmpfile('.png')
    run('qauto %s %s' % (im, out))
    return out


def points_apply_homography(H, pts):
    """
    Applies an homography to a list of 2D points.

    Args:
        H: numpy array containing the 3x3 homography matrix
        pts: numpy array containing the list of 2D points, one per line

    Returns:
        a numpy array containing the list of transformed points, one per line
    """
    # if the list of points is not a numpy array, convert it
    if (type(pts) == list):
        pts = np.array(pts)

    # convert the input points to homogeneous coordinates
    if len(pts[0]) < 2:
        print("""points_apply_homography: ERROR the input must be a numpy array
          of 2D points, one point per line""")
        return
    pts = np.hstack((pts[:, 0:2], pts[:, 0:1]*0+1))

    # apply the transformation
    Hpts = (np.dot(H, pts.T)).T

    # normalize the homogeneous result and trim the extra dimension
    Hpts = Hpts * (1.0 / np.tile( Hpts[:, 2], (3, 1)) ).T
    return Hpts[:, 0:2]


def bounding_box2D(pts):
    """
    bounding box for the points pts
    """
    dim = len(pts[0])  # should be 2
    bb_min = [min([t[i] for t in pts]) for i in range(dim)]
    bb_max = [max([t[i] for t in pts]) for i in range(dim)]
    return bb_min[0], bb_min[1], bb_max[0] - bb_min[0], bb_max[1] - bb_min[1]


def image_crop_gdal(im, x, y, w, h, out=None):
    """
    Crop an image using gdal_translate.

    Args:
        im: path to an image file
        x, y, w, h: four integers definig the rectangular crop in the image.
            (x, y) is the top-left corner, and (w, h) are the dimensions of the
            rectangle.
        out (optional): path to the output crop

    Returns:
        path to cropped tif image
    """
    if int(x) != x or int(y) != y:
        print('WARNING: image_crop_gdal will round the coordinates of your crop')

    if out is None:
        out = tmpfile('.tif')

    # do the crop with gdal_translate
    run(('gdal_translate -ot Float32 -co TILED=YES -co BIGTIFF=IF_NEEDED '
         '-srcwin %d %d %d %d %s %s') % (x, y, w, h, shellquote(im),
                                         shellquote(out)))

    return out


def run_binary_on_list_of_points(points, binary, option=None, env_var=None):
    """
    Runs a binary that reads its input on stdin.

    Args:
        points: numpy array containing all the input points, one per line
        binary: path to the binary. It is supposed to write one output value on
            stdout for each input point
        option: optional option to pass to the binary
        env_var (optional): environment variable that modifies the behaviour of
            the binary. It is a tuple containing 2 strings, eg ('PATH', '/bin')

    Returns:
        a numpy array containing all the output points, one per line.
    """
    # send the input points to stdin
    pts_file = tmpfile('.txt')
    np.savetxt(pts_file, points, '%.18f')
    p1 = subprocess.Popen(['cat', pts_file], stdout=subprocess.PIPE)

    # run the binary
    env = os.environ.copy()
    if env_var is not None:
        env[env_var[0]] = env_var[1]
    cmd = [binary]
    if option is not None:
        cmd.append(option)
    p2 = subprocess.Popen(cmd, env=env, stdin=p1.stdout,
                          stdout=subprocess.PIPE)

    # recover output values
    out = []
    for i in range(len(points)):
        out.append([float(x) for x in p2.stdout.readline().split()])

    return np.array(out)


def get_rectangle_coordinates(im):
    """
    Get the coordinates of a rectangle defined by the user's clicks.

    Args:
        im: path to an image to be displayed.

    Returns:
        x, y, w, h: coordinates of the rectangle selected by the user. x, y are the
            coordinates of the top-left corner, while (w, h) is the size of the
            rectangle.
    """
    points_file = tmpfile('.txt')
    run('python s2p/viewGL.py %s > %s' % (shellquote(im), points_file))
    x1, y1, x2, y2 = map(int, open(points_file).read().split())
    # viewGL.py returns the coordinates of two corners defining the rectangle.
    # We can's make any assumption on the ordering of these coordinates.

    x = min(x1, x2)
    w = max(x1, x2) - x
    y = min(y1, y2)
    h = max(y1, y2) - y
    return x, y, w, h


def get_roi_coordinates(img, preview):
    """
    Coordinates of a rectangle in a large image from user clicks on a preview.

    Args:
        img: path to the large image file
        preview: path to the preview image file

    Returns:
        x, y, w, h: coordinates of the rectangle selected by the user, in the
            large image frame. x, y are the coordinates of the top-left corner,
            while (w, h) is the size of the rectangle.

    A preview image is displayed, on which the user selects a rectangle.
    """
    # read preview/full images dimensions
    nc, nr = image_size_gdal(img)[:2]
    nc_preview, nr_preview = image_size_gdal(preview)[:2]

    # get the rectangle coordinates
    x, y, w, h = get_rectangle_coordinates(preview)

    # rescale according to preview/full ratio
    x = int(x*nc/nc_preview)
    y = int(y*nr/nr_preview)
    w = int(w*nc/nc_preview)
    h = int(h*nr/nr_preview)
    return x, y, w, h


def cargarse_basura(inputf, outputf):
    se=5
    tmp1 = outputf + '1.tif'
    tmp2 = outputf + '2.tif'
    tmpM = outputf + 'M.tif'
    run('morphoop %s min %d %s' % (inputf, se, tmpM))
    run('morphoop %s max %d %s' % (inputf, se, tmp1))
    run('morphoop %s max %d %s' % (inputf, se, tmpM))
    run('morphoop %s min %d %s' % (inputf, se, tmp2))
    run('plambda %s %s %s "x y - fabs %d > nan z if" -o %s' % (tmp1, tmp2, inputf, 5, tmpM))
    run('remove_small_cc %s %s %d %d' % (tmpM, outputf, 200, 5))
    run('rm -f %s %s %s' % (tmp1, tmp2, tmpM))


def print_elapsed_time(since_first_call=False):
    """
    Print the elapsed time since the last call or since the first call.

    Args:
        since_first_call:
    """
    t2 = datetime.datetime.now()
    if since_first_call:
        print("Total elapsed time:", t2 - print_elapsed_time.t0)
    else:
        try:
            print("Elapsed time:", t2 - print_elapsed_time.t1)
        except AttributeError:
            print("Elapsed time:", t2 - print_elapsed_time.t0)
    print_elapsed_time.t1 = t2
    print()
