# 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>


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


def run(cmd, env=os.environ, timeout=None):
    """
    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
        timeout (optional, int): time in seconds after which the function will
            raise an error if the command hasn't returned

    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()
    subprocess.check_call(cmd, shell=True, stdout=sys.stdout,
                          stderr=sys.stderr, env=env, timeout=timeout)
    print(datetime.datetime.now() - t)


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 crop_array(img, x, y, w, h, fill_value=0):
    """
    Crop an image represented as an array.

    Args:
        img (array): 2D input image
        x, y (ints): coordinate of the top-left corner of the crop
        w, h (ints): width and height of the crop
        fill_value (img.dtype): constant value used for filling the crop
            outside the input image domain

    Returns:
        array with the cropped image
    """
    crop = fill_value * np.ones((h, w), dtype=img.dtype)

    y0 = max(y, 0)
    y1 = min(y + h, img.shape[0])
    x0 = max(x, 0)
    x1 = min(x + w, img.shape[1])

    if y0 < y1 and x0 < x1:  # the requested crop overlaps the image
        crop[y0 - y:y1 - y, x0 - x:x1 - x] = img[y0:y1, x0:x1]

    return crop


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()
