from __future__ import division
from builtins import str
from builtins import range
from past.utils import old_div
import numpy as np
from osgeo import gdal, gdal_array
from girs.rast.parameter import RasterParameters
from girs.rast.raster import get_driver, RasterReader, RasterWriter
from girs.geom.envelope import merge_envelopes, buffer_envelope


def get_parameters_from_raster(r):
    try:
        r = RasterReader(r)
    except AttributeError:
        pass
    try:
        return r.get_parameters()
    except AttributeError:
        return None


def mosaic(*args, **kwargs):
    """

    :param args:
    :param kwargs:
    :return:
    """
    # Handle the case of *args being a tuple of list
    rp0 = get_parameters_from_raster(args[0])
    if not rp0:
        args = args[0]
    rp0 = get_parameters_from_raster(args[0])
    gt0 = rp0.geo_trans

    raster_parameters = [rp0]
    for i in range(1, len(args)):
        raster_parameters.append(get_parameters_from_raster(args[i]))
    envs = [rp0.get_extent_world() for rp0 in raster_parameters]
    env = merge_envelopes(envs, extent_type='union')
    env = buffer_envelope(env, (-gt0[1]*0.5, gt0[5]*0.5))
    (u_min, u_max, v_min, v_max), geo_trans = rp0.extent_world_to_pixel(*env)
    u_min, u_max = 0, u_max - u_min
    v_min, v_max = 0, v_max - v_min
    rp_out = RasterParameters(u_max+1, v_max+1, geo_trans, rp0.srs, rp0.number_of_bands, rp0.nodata, rp0.data_types,
                              driver_short_name=rp0.driverShortName)
    r_out = RasterWriter(rp_out, kwargs.pop('output_raster', None))
    for i in range(rp_out.number_of_bands):
        bn = i + 1
        dtype = gdal_array.GDALTypeCodeToNumericTypeCode(rp_out.data_types[i])
        array_out = np.full((rp_out.RasterYSize, rp_out.RasterXSize), rp_out.nodata[i], dtype)
        for arg in args:
            try:
                r = RasterReader(arg)
            except AttributeError:
                r = arg
            p = r.get_parameters()
            x_min, y_max = p.pixel_to_world(0, 0)
            x_min, y_max = x_min + (p.geo_trans[1] * 0.5), (y_max + p.geo_trans[5] * 0.5)
            u_min, v_min = rp_out.world_to_pixel(x_min, y_max)
            array = r.get_array(band_number=bn)
            array_out[v_min:v_min+p.RasterYSize, u_min:u_min+p.RasterXSize] = array
        r_out.set_array(array=array_out, band_number=bn)
    return r_out


def composite(*args, **kwargs):
    """

    :param args: Raster or filenames
    :param kwargs:
        :key output_raster: filename of the output raster
    :return:
    """
    p = get_parameters_from_raster(args[0])
    if not p:
        args = args[0]
    p = get_parameters_from_raster(args[0])
    p.number_of_bands = len(args)
    r_out = RasterWriter(p, kwargs.pop('output_raster', None))
    for i, arg in enumerate(args):
        try:
            r = RasterReader(arg)
        except AttributeError:
            r = arg
        r_out.set_array(r.get_array(), i+1)
    r_out.dataset.FlushCache()
    return r_out


def calculate(calc, **kwargs):
    """Examples::

        # Returns in 'MEM', applying the algebra to raster R, bands 1 and 2
        calculate("(R1-R2)/(R1+R2)", R='C:/tmp/raster1.tif')
        # Save to file
        calculate("(R1-R2)/(R1+R2)", R='C:/tmp/raster1.tif', output_raster='C:/tmp/ndvi.tif')
        # Using lists of bands. Note that [1:2] means band 1 and band 2. band numbers start at
        # 1 and include the last index
        calculate("(R[1, 2]+S[1, 2])/(R[1, 2]-S[1, 2])")
        # ... which is equivalent to
        calculate("(R[:2]+S[:2])/(R[:2]-S[:2])")
        # Applying calculate to all bands
        calculate("(R[:]+S[:])/(R[:]-S[:])")
        # This is also possible:
        calculate("(R[1, 2]+S[3, 4])/(R[:2]-S[1, 4])")

    See also:

        - https://svn.osgeo.org/gdal/trunk/gdal/swig/python/scripts/gdal_calc.py
        - https://github.com/rveciana/geoexamples/blob/master/python/gdal-performance/classification_blocks.py

    :param calc: algebraic operation supported by numpy arrays "(R1-R2)/(R1+R2)", where R1 and R2 are the keys used in the argument rasters:

        - R1: raster R band 1
        - R2: raster R band 2

    :param kwargs:

        - `<letter>` (a letter used in calc, string):  filename or Raster instance
        - `output_raster` (str): filename of the new rasters
        - `driver` (str): driver name

    :return:
    """
    output_raster = kwargs.pop('output_raster', '')
    driver = kwargs.pop('driver', None)
    driver = get_driver(output_raster, driver)
    driver = driver.ShortName

    r = RasterReader(list(kwargs.values())[0])

    # Separate list of bands in single calculations
    calc = _parse_range_to_list(calc, r.get_band_count())
    calculations = _split_calculation(calc)

    raster_parameters = r.get_parameters()
    raster_parameters.driverShortName = driver
    raster_parameters.number_of_bands = len(calculations)

    r_out = RasterWriter(raster_parameters, output_raster)

    for i, calc in enumerate(calculations):
        calc_dict = _get_symbol_and_masked_arrays(calc, kwargs)
        arr = eval(calc, calc_dict)
        r_out.set_array(arr.data, i+1)
    r_out.dataset.FlushCache()
    return r_out


def resample(input_raster, pixel_sizes, resample_alg=gdal.GRA_NearestNeighbour, **kwargs):
    """
    Args:
        input_raster: file name of a rasters or a Raster object

        pixel_sizes (list, str or Raster): list of pixel_width, pixel_height, file name of a rasters, or Raster object

        kwargs:
            'output_raster' (string): filename of the new rasters.
            'driver' (string): driver name (e.g., 'MEM')

    """
    try:
        pixel_width, pixel_height = float(pixel_sizes), float(pixel_sizes)
    except:
        try:
            pixel_width, pixel_height = pixel_sizes
            pixel_width, pixel_height = float(pixel_width), float(pixel_height)
        except:
            try:
                raster0 = RasterReader(pixel_sizes)
            except:
                raster0 = pixel_sizes
            pixel_width, pixel_height = raster0.get_pixel_size()
            del raster0

    try:
        input_raster = RasterReader(input_raster)
    except:
        pass

    x_min, x_max, y_min, y_max = input_raster.get_extent()

    driver = get_driver(kwargs['output_raster'] if 'output_raster' in kwargs else None,
                        kwargs['driver'] if 'driver' in kwargs else None)
    driver = driver.ShortName

    len_x = x_max-x_min
    len_y = y_max-y_min
    u_max, v_max = int(old_div(len_x,pixel_width)), int(old_div(len_y,pixel_height))
    if (old_div((len_x - u_max * pixel_width), pixel_width)) > 0.5:
        u_max += 1
    if (old_div((len_y - v_max * pixel_height), pixel_height)) > 0.5:
        v_max += 1

    gt_out = list(input_raster.get_geotransform())  # make it mutable
    gt_out[1], gt_out[5] = pixel_width, -pixel_height

    name = kwargs.get('output_raster', None)
    driver = kwargs.get('driver', None)
    raster_parameters = input_raster.get_parameters()
    raster_parameters.RasterXSize = u_max
    raster_parameters.RasterYSize = v_max
    raster_parameters.geo_trans = gt_out
    raster_parameters.driverShortName = driver
    raster_out = RasterWriter(raster_parameters, name)

    gdal.ReprojectImage(input_raster.dataset, raster_out.dataset, input_raster.get_coordinate_system(), raster_out.get_coordinate_system(), resample_alg)

    return raster_out


def strip(input_raster):
    """Remove top and bottom rows, and left and right columns containing only nodata

    :param input_raster:
    :return:
    """
    import numpy as np
    raster_parameters = input_raster.get_parameters()
    u_min = v_min = 0
    u_max, v_max = input_raster.get_raster_size()
    u_min_min, v_min_min = u_max, v_max
    u_max_max, v_max_max = u_min, v_min
    for bn in range(1, input_raster.get_band_count() + 1):
        array = input_raster.get_array()
        nodata = input_raster.get_nodata(bn)
        if nodata is np.nan or nodata is None:
            for u_min in range(u_min, u_max):
                if not np.all(np.isnan(array[:, u_min])):
                    break
            for u_max in range(u_max-1, u_min, -1):