# -*- coding: utf-8 -*-
#
#  xlsm.py
#  pangaea
#
#  Author : Alan D Snow, 2017.
#  License: BSD 3-Clause
"""pangea.xlsm
    This module is an extension for xarray for land surface models.
    (see: http://xarray.pydata.org/en/stable/internals.html#extending-xarray)
"""
from affine import Affine
import numpy as np
from osgeo import osr, gdalconst
import pandas as pd
from pyproj import Proj, transform
import wrf
import xarray as xr

from sloot.grid import (geotransform_from_yx, gdal_reproject,
                        resample_grid, utm_proj_from_latlon,
                        ArrayGrid)


@xr.register_dataset_accessor('lsm')
class LSMGridReader(object):
    """
    This is an extension for xarray specifically
    designed for land surface models.

    Example::

        import xarray as xr

        with xr.open_dataset('/path/to/file.nc') as xds:
            print(xds.lsm.projection)
    """
    def __init__(self, xarray_obj):
        self._obj = xarray_obj
        self._projection = None
        self._epsg = None
        self._geotransform = None
        self._affine = None
        self._center = None

        # set variable information
        self.y_var = 'lat'
        self.x_var = 'lon'
        self.time_var = 'time'
        # set dimension information
        self.y_dim = 'y'
        self.x_dim = 'x'
        self.time_dim = 'time'
        # convert lon from [0 to 360] to [-180 to 180]
        self.lon_to_180 = False

    def to_datetime(self):
        """Converts time to datetime."""
        time_values = self._obj[self.time_var].values
        if 'datetime' not in str(time_values.dtype):
            try:
                time_values = [time_val.decode('utf-8') for
                               time_val in time_values]
            except AttributeError:
                pass

            try:
                datetime_values = pd.to_datetime(time_values)
            except ValueError:
                # WRF DATETIME FORMAT
                datetime_values = \
                    pd.to_datetime(time_values,
                                   format="%Y-%m-%d_%H:%M:%S")

            self._obj[self.time_var].values = datetime_values

    @property
    def datetime(self):
        """Get datetime object for time variable"""
        self.to_datetime()
        return pd.to_datetime(self._obj[self.time_var].values)

    def _load_wrf_projection(self):
        """Load the osgeo.osr projection for WRF Grid.

        - 'MAP_PROJ': The map projection type as an integer.
        - 'TRUELAT1': True latitude 1.
        - 'TRUELAT2': True latitude 2.
        - 'MOAD_CEN_LAT': Mother of all domains center latitude.
        - 'STAND_LON': Standard longitude.
        - 'POLE_LAT': Pole latitude.
        - 'POLE_LON': Pole longitude.
        """
        # load in params from WRF Global Attributes
        possible_proj_params = ('MAP_PROJ', 'TRUELAT1', 'TRUELAT2',
                                'MOAD_CEN_LAT', 'STAND_LON', 'POLE_LAT',
                                'POLE_LON', 'CEN_LAT', 'CEN_LON', 'DX', 'DY')
        proj_params = dict()
        for proj_param in possible_proj_params:
            if proj_param in self._obj.attrs:
                proj_params[proj_param] = self._obj.attrs[proj_param]

        # determine projection from WRF Grid
        proj = wrf.projection.getproj(**proj_params)

        # export to Proj4 and add as osr projection
        self._projection = osr.SpatialReference()
        self._projection.ImportFromProj4(str(proj.proj4()))

    def _load_grib_projection(self):
        """Get the osgeo.osr projection for Grib Grid.
            - grid_type:  Lambert Conformal
            - Latin1:     True latitude 1.
            - Latin2:     True latitude 2.
            - Lov:        Central meridian.
            - Lo1:        Pole longitude.
            - La1:        Pole latitude.
            - Dx:         [ 3.]
            - Dy:         [ 3.]
        """
        lat_var_attrs = self._obj[self.y_var].attrs
        if 'Lambert Conformal' in lat_var_attrs['grid_type']:
            mean_lat = self._obj[self.y_var].mean().values
            proj4_str = ("+proj=lcc "
                         "+lat_1={true_lat_1} "
                         "+lat_2={true_lat_2} "
                         "+lat_0={latitude_of_origin} "
                         "+lon_0={central_meridian} "
                         "+x_0=0 +y_0=0 "
                         "+ellps=WGS84 +datum=WGS84 "
                         "+units=m +no_defs") \
                .format(true_lat_1=lat_var_attrs['Latin1'][0],
                        true_lat_2=lat_var_attrs['Latin2'][0],
                        latitude_of_origin=mean_lat,
                        central_meridian=lat_var_attrs['Lov'][0])
        else:
            raise ValueError("Unsupported projection: {grid_type}"
                             .format(grid_type=lat_var_attrs['grid_type']))

        # export to Proj4 and add as osr projection
        self._projection = osr.SpatialReference()
        self._projection.ImportFromProj4(proj4_str)

    @property
    def projection(self):
        """:func:`osgeo.osr.SpatialReference`
            The projection for the dataset.
        """
        if self._projection is None:
            # read projection information from global attributes
            map_proj4 = self._obj.attrs.get('proj4')
            if map_proj4 is not None:
                self._projection = osr.SpatialReference()
                self._projection.ImportFromProj4(str(map_proj4))
            elif 'MAP_PROJ' in self._obj.attrs:
                self._load_wrf_projection()
            elif 'grid_type' in self._obj[self.y_var].attrs:
                self._load_grib_projection()
            else:
                # default to EPSG 4326
                self._projection = osr.SpatialReference()
                self._projection.ImportFromEPSG(4326)
            # make sure EPSG loaded if possible
            self._projection.AutoIdentifyEPSG()
        return self._projection

    @property
    def epsg(self):
        """str: EPSG code"""
        if self._epsg is None:
            self._epsg = self.projection.GetAuthorityCode(None)
        return self._epsg

    @property
    def dx(self):
        """float: Pixel size in x direction."""
        return self.geotransform[1]

    @property
    def dy(self):
        """float: Pixel size in y direction."""
        return -self.geotransform[-1]

    @property
    def geotransform(self):
        """:obj:`tuple`: The geotransform for grid."""
        if self._geotransform is None:
            if self._obj.attrs.get('geotransform') is not None:
                self._geotransform = [float(g) for g in
                                      self._obj.attrs.get('geotransform')]

            elif str(self.epsg) != '4326':
                proj_y, proj_x = self.coords
                self._geotransform = geotransform_from_yx(proj_y,
                                                          proj_x)
            else:
                self._geotransform = geotransform_from_yx(*self.latlon)

        return self._geotransform

    @property
    def affine(self):
        """:func:`Affine`: The affine for the transformation."""
        if self._affine is None:
            self._affine = Affine.from_gdal(*self.geotransform)
        return self._affine

    @property
    def x_size(self):
        """int: Number of columns in the dataset."""
        return self._obj.dims[self.x_dim]

    @property
    def y_size(self):
        """int: Number of rows in the dataset."""
        return self._obj.dims[self.y_dim]

    @property
    def latlon(self):
        """Returns lat,lon arrays"""
        if 'MAP_PROJ' in self._obj.attrs:
            lat, lon = wrf.latlon_coords(self._obj, as_np=True)
        else:
            lon = self._obj[self.x_var].values
            lat = self._obj[self.y_var].values

        if lat.ndim == 3:
            lat = lat[0]
        if lon.ndim == 3:
            lon = lon[0]

        if 'MAP_PROJ' in self._obj.attrs:
            # WRF Grid is upside down
            lat = lat[::-1]
            lon = lon[::-1]

        if self.lon_to_180:
            lon = (lon + 180) % 360 - 180  # convert [0, 360] to [-180, 180]

        if lat.ndim < 2:
            lon, lat = np.meshgrid(lon, lat)

        return lat, lon

    @property
    def coords(self):
        """Returns y, x coordinate arrays"""
        lat, lon = self.latlon
        x_coords, y_coords = transform(Proj(init='epsg:4326'),
                                       Proj(self.projection.ExportToProj4()),
                                       lon,
                                       lat)
        return y_coords, x_coords

    @property
    def center(self):
        """Return the geographic center point of this dataset."""
        if self._center is None:
            # we can use a cache on our accessor objects, because accessors
            # themselves are cached on instances that access them.
            lat, lon = self.latlon
            self._center = (float(np.nanmean(lon)), float(np.nanmean(lat)))
        return self._center

    def _export_dataset(self, variable, new_data, grid):
        """Export subset of dataset."""
        lats, lons = grid.latlon

        return xr.Dataset({variable: (['time', 'y', 'x'],
                                      new_data,
                                      self._obj[variable].attrs),
                           },
                          coords={'lat': (['y', 'x'],
                                          lats,
                                          self._obj[variable]
                                          .coords[self.y_var].attrs),
                                  'lon': (['y', 'x'],
                                          lons,
                                          self._obj[variable]
                                          .coords[self.x_var].attrs),
                                  'time': (['time'],
                                           self._obj[self.time_var].values,
                                           self._obj[self.time_var].attrs),
                                  },
                          attrs={'proj4': grid.proj4,
                                 'geotransform': grid.geotransform,
                                 }
                          )

    def resample(self, variable, match_grid):
        """Resample data to grid.

            Parameters
            ----------
            variable: :obj:`str`
                Name of variable in dataset.
            match_grid: :func:`gdal.Dataset` or :func:`sloot.grid.GDALGrid`
                Grid you want the data resampled to match resolution.
                You can also pass the path to the grid.
        """
        new_data = []
        for band in range(self._obj.dims[self.time_dim]):
            data = self._obj[variable][band].values
            arr_grid = ArrayGrid(in_array=data,
                                 wkt_projection=self.projection.ExportToWkt(),
                                 geotransform=self.geotransform)
            resampled_data_grid = resample_grid(original_grid=arr_grid,
                                                match_grid=match_grid,
                                                as_gdal_grid=True)
            new_data.append(resampled_data_grid.np_array())

        self.to_datetime()
        return self._export_dataset(variable, np.array(new_data),
                                    resampled_data_grid)

    def _getvar(self, variable, yslice, xslice):
        """Get the variable either directly or calculated"""
        # FAILED ATTEMPT TO USE wrf.getvar
        # if 'MAP_PROJ' in self._obj.attrs:
        #    try:
        #        nc_file = self._obj._file_obj.ds
        #    except AttributeError:
        #        nc_file = self._obj._file_obj.file_objs
        #    var = wrf.getvar(nc_file, variable)
        def extract_slice(var, slice_arr):
            """extract by slice"""
            if var.ndim == 4:
                var = var[slice_arr[0], slice_arr[1],
                          slice_arr[2], slice_arr[3]]
            if var.ndim == 3:
                var = var[slice_arr[0], slice_arr[1], slice_arr[2]]
            else:
                var = var[slice_arr[0], slice_arr[1]]
            return var

        var = self._obj[variable]
        slc = [slice(None)] * var.ndim
        # flip in y-direction
        if 'MAP_PROJ' in self._obj.attrs:
            # WRF Y DIM IS BACKWARDS
            slc[var.get_axis_num(self.y_dim)] = slice(None, None, -1)
            var = extract_slice(var, slc)
        # get data out
        slc[var.get_axis_num(self.x_dim)] = xslice
        slc[var.get_axis_num(self.y_dim)] = yslice
        return extract_slice(var, slc)

    def getvar(self, variable,
               yslice=slice(None),
               xslice=slice(None),
               calc_4d_method=None,
               calc_4d_dim=None):
        """Get variable from model with subset options.

            Parameters
            ----------
            variable: :obj:`str`
                Name of variable in dataset.
            yslice: :obj:`slice`, optional
                Slice in y-direction of grid to extract data from.
            xslice: :obj:`slice`, optional
                Slice in x-direction of grid to extract data from.
            calc_4d_method: :obj:`str`
                Method to convert 4D variables to 3D variables
                (Ex. 'mean', 'min', or 'max').
            calc_4d_dim: :obj:`str`
                Dimension to reduce grid from 4D to 3D (Ex. 'top_bottom').

            Returns
            -------
            :func:`xarray.DataArray`
        """
        data = self._getvar(variable, yslice, xslice)

        if data.ndim == 4:
            if calc_4d_method is None or calc_4d_dim is None:
                raise ValueError("The variable {var} has 4 dimension. "
                                 "Need 'calc_4d_method' and 'calc_4d_dim' "
                                 "to proceed ...".format(var=variable))
            data = getattr(data, calc_4d_method)(dim=calc_4d_dim)

        data[self.time_var] = self._obj[self.time_var]

        return data

    def to_projection(self, variable, projection):
        """Convert Grid to New Projection.

            Parameters
            ----------
            variable: :obj:`str`
                Name of variable in dataset.
            projection: :func:`osr.SpatialReference`
                Projection to convert data to.

            Returns
            -------
            :func:`xarray.Dataset`
        """
        new_data = []
        for band in range(self._obj.dims[self.time_dim]):
            arr_grid = ArrayGrid(in_array=self._obj[variable][band].values,
                                 wkt_projection=self.projection.ExportToWkt(),
                                 geotransform=self.geotransform)
            ggrid = gdal_reproject(arr_grid.dataset,
                                   src_srs=self.projection,
                                   dst_srs=projection,
                                   resampling=gdalconst.GRA_Average,
                                   as_gdal_grid=True)
            new_data.append(ggrid.np_array())

        self.to_datetime()
        return self._export_dataset(variable, np.array(new_data),
                                    ggrid)

    def to_utm(self, variable):
        """Convert Grid to UTM projection at center of grid.

            Parameters
            ----------
            variable: :obj:`str`
                Name of variable in dataset.

            Returns
            -------
            :func:`xarray.Dataset`
        """
        # get utm projection
        center_lon, center_lat = self.center
        dst_proj = utm_proj_from_latlon(center_lat, center_lon,
                                        as_osr=True)
        return self.to_projection(variable, dst_proj)

    def to_tif(self, variable, time_index, out_path):
        """Dump a variable at a time index to a geotiff.

            Parameters
            ----------
            variable: :obj:`str`
                Name of variable in dataset.
            time_index: int
                0-based time index,
            out_path: :obj:`str`
                Path to output geotiff file,
        """
        arr_grid = ArrayGrid(in_array=self._obj[variable][time_index].values,
                             wkt_projection=self.projection.ExportToWkt(),
                             geotransform=self.geotransform)
        arr_grid.to_tif(out_path)
