"""
   Copyright 2015 Creare

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

from numpy.linalg import inv
from traits.api import *
import numpy as np
import gdal
import gdalconst
import matplotlib, matplotlib.cm

NO_DATA_VALUE = -9999

d_name_to_wkt = {'WGS84' : r'GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433],AUTHORITY["EPSG","4326"]]',
                 'NAD83' : r'GEOGCS["NAD83",DATUM["North_American_Datum_1983",SPHEROID["GRS 1980",6378137,298.2572221010002,AUTHORITY["EPSG","7019"]],TOWGS84[0,0,0,0,0,0,0],AUTHORITY["EPSG","6269"]],PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433],AUTHORITY["EPSG","4269"]]',
                }
d_name_to_epsg = {'WGS84' : 4326,
                  'NAD83': 4269
                 }
d_wkt_to_name = {v:k for k, v in d_name_to_wkt.iteritems()}
d_wkt_to_name[r'GEOGCS["NAD83",DATUM["North_American_Datum_1983",SPHEROID["GRS 1980",6378137,298.2572221010002,AUTHORITY["EPSG","7019"]],AUTHORITY["EPSG","6269"]],PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433],AUTHORITY["EPSG","4269"]]'] = 'NAD83'
d_wkt_to_name[r'GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433],AUTHORITY["EPSG","4326"]]'] = 'WGS84'  # afghanistan dem

d_epsg_to_name = {4326: 'WGS84',
                  4269: 'NAD83',
                 }

# This trait maps a user-friendly name (e.g., WGS84) to an official WKT.
projection_wkt_trait = Trait('WGS84',
             d_name_to_wkt)

class Point(HasStrictTraits):
    """
    Point(Lat, Lon)

    A simple convenience class to deal with latitude and longitude, which
    can be error-prone to convert otherwise.

    Examples
    --------
    >>> Point(-1.426667, -20.583611)
    <Lat,Lon: -1.427, -20.584>

    >>> Point((1, 25, 36), (20, 35, 1), name="Site A")
    <Site A: 1.427, 20.584>

    >>> p = Point(-1.426667, -20.583611)
    >>> p.lat
    -1.426667
    >>> p.lat_dms
    (1.0, 25.0, 36.001199999999599)
    >>> p.lat_dms = (1.0, 2.0, 3.0)
    >>> p
    <Lat,Lon: 1.034, -20.584>

    Conventions
    -----------
    Following ISO 6709,
        * LAT, LON
        * North (lat) and East (lon) are positive.
        * South (lat) and West (lon) are negative.
        * Decimal representation is preferred; Sexagesimal (base-60) is allowed.

    """
    name = Str()
    wkt = projection_wkt_trait

    lon = Float()
    lon_dms = Property(Tuple((Float, Float, Float)))
    def _set_lon_dms(self, dms):
        deg, min, sec = dms
        self.lon = np.sign(deg) * (abs(deg) + min / 60.0 + sec / 3600.0)
    def _get_lon_dms(self):
        deg = np.floor(abs(self.lon))
        min = np.floor((abs(self.lon) - deg) * 60.0)
        sec = np.round((abs(self.lon) - deg - min / 60.0) * 3600.0,
                       4)  # round to 4 decimal places.
        return (np.sign(self.lon) * deg, min, sec)

    lat = Float()
    lat_dms = Property(Tuple((Float, Float, Float)))
    def _set_lat_dms(self, dms):
        deg, min, sec = dms
        self.lat = np.sign(deg) * (abs(deg) + min / 60.0 + sec / 3600.0)
    def _get_lat_dms(self):
        deg = np.floor(abs(self.lat))
        min = np.floor((abs(self.lat) - deg) * 60.0)
        sec = np.round((abs(self.lat) - deg - min / 60.0) * 3600.0,
                       4)  # round to 4 decimal places.
        return (np.sign(self.lat) * deg, min, sec)

    def to_wkt(self, target_wkt):
        # If we're going from WGS84 -> Spherical Mercator, use PyProj, because
        #  there seems to be a bug in OGR that gives us an offset. (GDAL
        #  does fine, though.
        if target_wkt == self.wkt:
            return self
        import osr
        dstSpatialRef = osr.SpatialReference()
        dstSpatialRef.ImportFromEPSG(d_name_to_epsg[target_wkt])
        # dstSpatialRef.ImportFromWkt(d_name_to_wkt[target_wkt])
        srcSpatialRef = osr.SpatialReference()
        srcSpatialRef.ImportFromEPSG(d_name_to_epsg[self.wkt])
        # srcSpatialRef.ImportFromWkt(self.wkt_)
        coordTransform = osr.CoordinateTransformation(srcSpatialRef, dstSpatialRef)
        a, b, c = coordTransform.TransformPoint(self.lon, self.lat)
        return Point(b, a, wkt=target_wkt)

    def __str__(self):
        return "<%s (%s): %02.3f, %03.3f>" % (self.name if self.name else 'Lat,Lon',
                                              self.wkt,
                                         self.lat, self.lon)

    def __repr__(self):
        s =  ("Point(%02.3f, %02.3f, wkt='%s'" %
              (self.lat, self.lon, self.wkt))
        if self.name:
            s += ", name='%s'" % self.name
        s += ")"

        return s

    def __init__(self, lat=None, lon=None, **kwargs):
        HasStrictTraits.__init__(self, **kwargs)
        if lon is not None:
            try:
                self.lon = lon  # float
            except:
                self.lon_dms = lon  # tuple
        if lat is not None:
            try:
                self.lat = lat  # float
            except:
                self.lat_dms = lat  # tuple


def grid_coords_from_corners(upper_left_corner, lower_right_corner, size):
    ''' Points are the outer edges of the UL and LR pixels. Size is rows, columns.
    GC projection type is taken from Points. '''
    assert upper_left_corner.wkt == lower_right_corner.wkt
    geotransform = np.array([upper_left_corner.lon, -(upper_left_corner.lon - lower_right_corner.lon) / float(size[1]), 0,
                            upper_left_corner.lat, 0, -(upper_left_corner.lat - lower_right_corner.lat) / float(size[0])])
    return GridCoordinates(geotransform=geotransform,
                               wkt=upper_left_corner.wkt,
                               y_size=size[0],
                               x_size=size[1])

class GridCoordinates(HasStrictTraits):
    """
    Defines mapping of input layers to real-world time and space.
    """
    date = Date()
    time = Time()
    geotransform = Array('float', [6],)
    wkt = projection_wkt_trait
    x_size = Int()
    y_size = Int()
    x_axis = Property(Array(), depends_on='geotransform, x_size')
    y_axis = Property(Array(), depends_on='geotransform, y_size')

    ULC = Property(Instance(Point), depends_on='x_axis, y_axis')
    URC = Property(Instance(Point), depends_on='x_axis, y_axis')
    LLC = Property(Instance(Point), depends_on='x_axis, y_axis')
    LRC = Property(Instance(Point), depends_on='x_axis, y_axis')
    def _get_ULC(self): return Point(self.geotransform[3], self.geotransform[0], wkt=self.wkt)
    def _get_URC(self): return Point(self.geotransform[3], self.geotransform[0] + self.geotransform[1] * self.x_size, wkt=self.wkt)
    def _get_LLC(self): return Point(self.geotransform[3] + self.geotransform[5] * self.y_size, self.geotransform[0], wkt=self.wkt)
    def _get_LRC(self): return Point(self.geotransform[3] + self.geotransform[5] * self.y_size, self.geotransform[0] + self.geotransform[1] * self.x_size, wkt=self.wkt)

    projection_wkt = Property(Str)  # For backwards compatibility w/ previously pickled layers.
    def _set_projection_wkt(self, val):
        self.wkt = data.layers.types.d_wkt_to_name[val]

    def __repr__(self):
        return '<GridCoordinates: %s -> %s, %d x %d>' % (
                         self.ULC, self.LRC, self.y_size, self.x_size)

    def intersects(self, other_grid_coordinates):
        """ returns True if the GC's overlap. """
        ogc = other_grid_coordinates  # alias
        # for explanation: http://stackoverflow.com/questions/306316/determine-if-two-rectangles-overlap-each-other
        # Note the flipped y-coord in this coord system.
        ax1, ay1, ax2, ay2 = self.ULC.lon, self.ULC.lat, self.LRC.lon, self.LRC.lat
        bx1, by1, bx2, by2 = ogc.ULC.lon, ogc.ULC.lat, ogc.LRC.lon, ogc.LRC.lat
        if ((ax1 <= bx2) and (ax2 >= bx1) and (ay1 >= by2) and (ay2 <= by1)):
            return True
        else:
            return False



    def unique_str(self):
        """ A string that (ideally) uniquely represents this GC object. This
        helps with naming files for caching. 'Unique' is defined as 'If
        GC1 != GC2, then GC1.unique_str() != GC2.unique_str()'; conversely,
        'If GC1 == GC2, then GC1.unique_str() == GC2.unique_str()'.

        The string should be filename-safe (no \/:*?"<>|).

        ..note::Because of length/readability restrictions, this fxn ignores
        wkt.

        Example output:
        "-180.000_0.250_0.000_90.000_0.000_-0.251_512_612_2013-05-21_12_32_52.945000"
        """

        unique_str = "_".join(["%.3f" % f for f in self.geotransform] +
                              ["%d" % d for d in self.x_size, self.y_size]
                              )
        if self.date is not None:
            unique_str += '_' + str(self.date)
        if self.time is not None:
            unique_str += '_' + str(self.time)
        return unique_str.replace(':', '_')

    def __eq__(self, other):
        return (isinstance(other, self.__class__)
            and np.allclose(self.geotransform, other.geotransform)
            and (self.x_size == other.x_size)
            and (self.y_size == other.y_size)
            and (self.date == other.date)
            and (self.wkt == other.wkt)
            and (self.time == other.time)
            )


    def __ne__(self, other):
        return not self.__eq__(other)

    @cached_property
    def _get_x_axis(self):
        """See http://www.gdal.org/gdal_datamodel.html for details."""
        # 0,0 is top/left top top/left pixel. Actual x/y coord of that pixel are (.5,.5).
        x_centers = np.linspace(.5, self.x_size - .5, self.x_size)
        y_centers = x_centers * 0
        return (self.geotransform[0]
                + self.geotransform[1] * x_centers
                + self.geotransform[2] * y_centers)

    @cached_property
    def _get_y_axis(self):
        """See http://www.gdal.org/gdal_datamodel.html for details."""
        # 0,0 is top/left top top/left pixel. Actual x/y coord of that pixel are (.5,.5).
        y_centers = np.linspace(.5, self.y_size - .5, self.y_size)
        x_centers = y_centers * 0
        return (self.geotransform[3]
                + self.geotransform[4] * x_centers
                + self.geotransform[5] * y_centers)



    def raster_to_projection_coords(self, pixel_x, pixel_y):
        """ Use pixel centers when appropriate.
        See documentation for the GDAL function GetGeoTransform for details. """
        h_px_py = np.array([1, pixel_x, pixel_y])
        gt = np.array([[1, 0, 0], self.geotransform[0:3], self.geotransform[3:6]])
        arr = np.inner(gt, h_px_py)
        return arr[2], arr[1]

    def projection_to_raster_coords(self, lat, lon):
        """ Returns pixel centers.
        See documentation for the GDAL function GetGeoTransform for details. """
        r_px_py = np.array([1, lon, lat])
        tg = inv(np.array([[1, 0, 0], self.geotransform[0:3], self.geotransform[3:6]]))
        return np.inner(tg, r_px_py)[1:]

    def _as_gdal_dataset(self, driver="MEM", n_raster_count=1,
                   file_name="memory.tif", data_type=gdalconst.GDT_Float32
                   ):
        driver = gdal.GetDriverByName(driver)
        dataset = driver.Create(file_name, int(self.x_size), int(self.y_size),
                                n_raster_count, data_type)

        dataset.SetGeoTransform(self.geotransform)  # dem_geotrans)
        dataset.SetProjection(self.wkt_)
        return dataset

    def copy_and_transform(self, zoom=1.0):
        # Probalby doesn't handle angled reference frames correctly.
        copy = self.clone_traits(copy='deep')

        x_center = self.geotransform[0] + self.geotransform[1] * self.x_size * .5
        x_new_spacing = self.geotransform[1] / zoom
        copy.geotransform[0:3] = x_center - x_new_spacing * self.x_size * .5, x_new_spacing, 0

        y_center = self.geotransform[3] + self.geotransform[5] * self.y_size * .5
        y_new_spacing = self.geotransform[5] / zoom
        copy.geotransform[3:6] = y_center - y_new_spacing * self.y_size * .5, 0, y_new_spacing

        return copy

class AbstractDataLayer(HasStrictTraits):

    #===========================================================================
    # Coordinate system maintained by parent simulation
    #===========================================================================
    grid_coordinates = Instance('GridCoordinates')

    #===========================================================================
    # Data Description
    #===========================================================================
    name = Str()  # e.g., 'Temperature'
    units = Str()  # e.g., 'K'
    data_type = Enum([None, 'LHTFL', 'SHTFL', 'GFLUX', 'SSRUN', 'BGRUN', 'TMP',
                 'ALBDO', 'WEASD', 'SNOD', 'SOLMST', 'TSOIL', 'SOLLIQ', 'EVAPTP',
                 'CANOPY', 'WDRUN', 'TMP', 'TMIN', 'TMAX', 'SPFH', 'PRES',
                 'DNSWF', 'DNLWF', 'LAND', 'VEGTYP', 'SOLTYP', 'TERR', 'VGREEN',
                 'RELSOL', 'MINRH', 'APCP', 'EVAPTS', 'NSWRS', 'NLWRS', 'SNOHF',
                 'SNOEV', 'DSWRF', 'DLWRF', 'ASNOW', 'ARAIN', 'EVP', 'SNOM',
                 'AVSFT', 'CNWAT', 'MSTAV', 'EVCW', 'TRANS', 'EVBS',
                 'SBSNO', 'PEVPR', 'ACOND', 'SNOWC', 'CCOND', 'RCS', 'RCT',
                  'RCQ', 'RCSOL', 'RSMIN', 'LAI', 'VEG', 'var250', 'var255'
                  ])
    depth = Enum([None, 'sfc', '0-200 cm down', '0-10 cm down', '0-100 cm down',
                  '0-200 cm down', '10-40 cm down',
                  '40-100 cm down', '100-200 cm down', '2 m above gnd',
                  '10 m above gnd'])

    def __repr__(self):
        return "<%s: %s; Data type: %s; Depth: %s >" % (self.__class__, self.name, self.data_type, self.depth)

    #===========================================================================
    # Enumeration
    #===========================================================================
    is_enumerated = Bool(False)
    enumeration_legend = Dict(key_trait=Int, value_trait=Str)  # e.g., {0:'Sea', 1:'Land'}
    enumeration_colors = Dict(key_trait=Int, value_trait=Tuple((1., 1., 1.)))  # e.g., {0:(128,128,128), 1:(1,50,150)}

    ############################################################################
    # Plotting Info (if not enumerated)
    ############################################################################
    scalar_c_lims = List()
    scalar_cm = Instance(matplotlib.colors.Colormap)

    def _scalar_cm_default(self):
        return matplotlib.cm.get_cmap('gray')

    def reproject_to_grid_coordinates(self, grid_coordinates, interp=gdalconst.GRA_NearestNeighbour):
        """ Reprojects data in this layer to match that in the GridCoordinates
        object. """
        source_dataset = self.grid_coordinates._as_gdal_dataset()
        dest_dataset = grid_coordinates._as_gdal_dataset()
        rb = source_dataset.GetRasterBand(1)
        rb.SetNoDataValue(NO_DATA_VALUE)
        rb.WriteArray(np.ma.filled(self.raster_data, NO_DATA_VALUE))

        gdal.ReprojectImage(source_dataset, dest_dataset,
                            source_dataset.GetProjection(),
                            dest_dataset.GetProjection(),
                            interp)
        dest_layer = self.clone_traits()
        dest_layer.grid_coordinates = grid_coordinates
        rb = dest_dataset.GetRasterBand(1)
        dest_layer.raster_data = np.ma.masked_values(rb.ReadAsArray(), NO_DATA_VALUE)
        return dest_layer

    def export_to_geotiff(self, file_name):
        dest_dataset = self.grid_coordinates._as_gdal_dataset(driver='GTiff',
                                                              file_name=file_name)
        rb = dest_dataset.GetRasterBand(1)
        rb.WriteArray(self.raster_data.filled())
        rb.SetNoDataValue(float(self.raster_data.fill_value))
        rb.SetDescription(self.name)
        rb.SetUnitType(self.units)

    def inpaint(self):
        """ Replace masked-out elements in an array using an iterative image inpainting algorithm. """

        import inpaint
        filled = inpaint.replace_nans(np.ma.filled(self.raster_data, np.NAN).astype(np.float32), 3, 0.01, 2)
        self.raster_data = np.ma.masked_invalid(filled)

    def colormap(self):
        from matplotlib import colors
        import collections
        if self.is_enumerated:
            if self.enumeration_colors:
                d = collections.OrderedDict(sorted(self.enumeration_colors.items()))
                cmap = colors.ListedColormap(d.values())  # + [(0., 0., 0.)])
                bounds = np.array(d.keys() + [d.keys()[-1] + 1]) - .5
                norm = colors.BoundaryNorm(bounds, cmap.N)
                return cmap, norm
            else:
                return None, None
        # not enumerated.
        return self.scalar_cm

    def to_rgba(self):
        data = self.raster_data
        if self.is_enumerated:
            if self.enumeration_colors:
                cmap, norm = self.colormap()
                data2 = norm(data)  # np.clip((data - MIN) / (MAX - MIN), 0, 1)
                rgba = (cmap(data2) * 255).astype(int)
                rgba[:, :, 3] = np.logical_not(data.mask).astype(int) * 255
                return rgba
            else:
                raise NotImplementedError()

        # Not enumerated...
        if self.scalar_c_lims:
            MIN, MAX = self.scalar_c_lims
        else:
            MIN, MAX = data.min(), data.max()

        cm = self.colormap()
        data2 = np.clip((data - MIN) / (MAX - MIN), 0, 1)
        rgba = (cm(data2) * 255).astype(int)
        rgba[:, :, 3] = np.logical_not(data.mask).astype(int) * 255

        return rgba


#===============================================================================
# DifferentiateInput/Staging/Results Layers
#===============================================================================
class InputLayerMixin(HasStrictTraits):
    duration_s = Float()

class StagingLayerMixin(HasStrictTraits):
    pass

class ResultsLayerMixin(HasStrictTraits):
    pass

#===============================================================================
# Raster Data
#===============================================================================
class AbstractRasterDataLayer(AbstractDataLayer):
    raster_data = Array()

    def interp_value(self, lat, lon, indexed=False):
        """ Lookup a pixel value in the raster data, performing linear interpolation
        if necessary. Indexed ==> nearest neighbor (*fast*). """
        (px, py) = self.grid_coordinates.projection_to_raster_coords(lat, lon)
        if indexed:
            return self.raster_data[round(py), round(px)]
        else:
#             from scipy.interpolate import interp2d
#             f_interp = interp2d(self.grid_coordinates.x_axis, self.grid_coordinates.y_axis, self.raster_data, bounds_error=True)
#             return f_interp(lon, lat)[0]
            from scipy.ndimage import map_coordinates
            ret = map_coordinates(self.raster_data, [[py], [px]], order=1)  # linear interp
            return ret[0]

class InputRasterDataLayer(InputLayerMixin, AbstractRasterDataLayer):
    pass

class StagingRasterDataLayer(StagingLayerMixin, AbstractRasterDataLayer):
    pass

class ResultsRasterDataLayer(ResultsLayerMixin, AbstractRasterDataLayer):
    pass

#===============================================================================
# Point Data
#===============================================================================
class PointMeasurement(HasStrictTraits):
    pass  # Uncertain how to define coord system and how this measurement looks...

class AbstractPointDataLayer(AbstractDataLayer):
    point_measurements = List(Instance(PointMeasurement))

class InputPointDataLayer(InputLayerMixin, AbstractPointDataLayer):
    pass

class StagingPointDataLayer(StagingLayerMixin, AbstractPointDataLayer):
    pass

class ResultsPointDataLayer(ResultsLayerMixin, AbstractPointDataLayer):
    pass


