"""Diagnostic variables for WRF output.

Diagnostic variables are simply a subclass of FakeVariable that implement
__getitem__. See examples below.
"""
from __future__ import division
import copy

import numpy as np
import pyproj
try:
    from pandas import to_datetime
    from xarray.core import indexing
except ImportError:
    pass

from salem import lazy_property, wgs84, gis


class Unstaggerer(object):
    """Duck NetCDF4.Variable class which "unstaggers" WRF variables.

     It looks for the staggered dimension and automatically unstaggers it.
     """

    def __init__(self, ncvar):
        """Instanciate.

        Parameters
        ----------
        ncvar: the netCDF variable to unstagger.
        """

        self.ncvar = ncvar

        # Attributes
        self.description = ncvar.description
        self.units = ncvar.units
        # Replace the dimension name
        dims = list(ncvar.dimensions)
        self.ds = np.nonzero(['_stag' in d for d in dims])[0][0]
        dims[self.ds] = dims[self.ds].replace('_stag', '')
        self.dimensions = dims
        shape = list(ncvar.shape)
        shape[self.ds] -=1
        self.shape = tuple(shape)

        # this is quickndirty, and probably wrong
        self.set_auto_maskandscale = ncvar.set_auto_maskandscale
        self.ncattrs = ncvar.ncattrs
        self.filters = ncvar.filters

        def _chunking():
            return self.shape
        self.chunking = _chunking

        for attr in self.ncattrs():
            setattr(self, attr, getattr(ncvar, attr))
        self.dtype = ncvar.dtype

        # Needed later
        self._ds_shape = ncvar.shape[self.ds]

    def getncattr(self, name):
        # dummy getncattrs
        return getattr(self, name)

    @staticmethod
    def can_do(ncvar):
        """Checks if the variable can be unstaggered.

        Parameters
        ----------
        ncvar: the netCDF variable candidate forunstagger.
        """
        return np.any(['_stag' in d for d in ncvar.dimensions])

    def __getitem__(self, item):
        """Override __getitem__."""

        # take care of ellipsis and other strange indexes
        item = list(indexing.expanded_indexer(item, len(self.dimensions)))

        # Slice to change
        sl = item[self.ds]
        if np.isscalar(sl) and not isinstance(sl, slice):
            sl = slice(sl, sl+1)

        # Ok, get the indexes right
        start = sl.start or 0
        stop = sl.stop or self._ds_shape
        if stop < 0:
            stop += self._ds_shape-1
        stop = np.clip(stop+1, 0, self._ds_shape)
        itemr = copy.deepcopy(item)
        item[self.ds] = slice(start, stop-1)
        itemr[self.ds] = slice(start+1, stop)
        return 0.5*(self.ncvar[item] + self.ncvar[itemr])


class FakeVariable(object):
    """Duck NetCDF4.Variable class
    """
    def __init__(self, nc):
        self.name = self.__class__.__name__
        self.nc = nc

    @staticmethod
    def can_do():
        raise NotImplementedError()

    def _copy_attrs_from(self, ncvar):
        # copies the necessary nc attributes from a template variable
        self.ncattrs = ncvar.ncattrs
        self.filters = ncvar.filters
        self.chunking = ncvar.chunking
        for attr in self.ncattrs():
            setattr(self, attr, getattr(ncvar, attr))
        self.dimensions = ncvar.dimensions
        self.dtype = ncvar.dtype
        self.set_auto_maskandscale = ncvar.set_auto_maskandscale
        self.shape = ncvar.shape

    def getncattr(self, name):
        # dummy getncattrs
        return getattr(self, name)

    def __getitem__(self, item):
        raise NotImplementedError()


class T2C(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['T2'])
        self.units = 'C'
        self.description = '2m Temperature'

    @staticmethod
    def can_do(nc):
        return 'T2' in nc.variables

    def __getitem__(self, item):
        return self.nc.variables['T2'][item] - 273.15


class AccumulatedVariable(FakeVariable):
    """Common logic for all accumulated variables."""

    def __init__(self, nc, accvn):
        FakeVariable.__init__(self, nc)
        self.accvn = accvn
        self._copy_attrs_from(nc.variables[self.accvn])
        # Needed later
        self._nel = nc.variables[self.accvn].shape[0]

    @lazy_property
    def _factor(self):
        # easy would be to have time step as variable
        vars = self.nc.variables
        if 'XTIME' in vars:
            dt_minutes = vars['XTIME'][1] - vars['XTIME'][0]
        else:
            # ok, parse time
            time = []
            stimes = vars['Times'][0:2]
            for t in stimes:
                time.append(to_datetime(t.tostring().decode(),
                                        errors='raise',
                                        format='%Y-%m-%d_%H:%M:%S'))
            dt_minutes = time[1] - time[0]
            dt_minutes = dt_minutes.seconds / 60

        return 60 / dt_minutes

    @staticmethod
    def can_do(nc):
        return nc.dimensions['Time'].size > 1

    def __getitem__(self, item):

        # take care of ellipsis and other strange indexes
        item = list(indexing.expanded_indexer(item, len(self.dimensions)))

        # time is always going to be first dim I hope
        sl = item[0]
        if np.isscalar(sl) and not isinstance(sl, slice):
            sl = slice(sl, sl+1)

        # Ok, get the indexes right
        start = sl.start or 0
        stop = sl.stop or self._nel
        if stop < 0:
            stop += self._nel-1
        start -= 1
        do_nan = False
        if start < 0:
            do_nan = True
        itemr = copy.deepcopy(item)
        item[0] = slice(start, stop-1)
        itemr[0] = slice(start+1, stop)

        # done
        var = self.nc.variables[self.accvn]
        if do_nan:
            item[0] = slice(0, stop-1)
            out = var[itemr]
            out[1:, ...] -= var[item]
            out[0, ...] = np.NaN
        else:
            out = var[itemr]
            out -= var[item]
        return out * self._factor


class PRCP_NC(AccumulatedVariable):

    def __init__(self, nc):
        AccumulatedVariable.__init__(self, nc, 'RAINNC')
        self.units = 'mm h-1'
        self.description = 'Precipitation rate from grid scale physics'

    @staticmethod
    def can_do(nc):
        return AccumulatedVariable.can_do(nc) and 'RAINNC' in nc.variables


class PRCP_C(AccumulatedVariable):

    def __init__(self, nc):
        AccumulatedVariable.__init__(self, nc, 'RAINC')
        self.units = 'mm h-1'
        self.description = 'Precipitation rate from cumulus physics'

    @staticmethod
    def can_do(nc):
        return AccumulatedVariable.can_do(nc) and 'RAINC' in nc.variables


class PRCP(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['RAINC'])
        self.units = 'mm h-1'
        self.description = 'Total precipitation rate'

    @staticmethod
    def can_do(nc):
        return AccumulatedVariable.can_do(nc) and \
               'RAINC' in nc.variables and 'RAINNC' in nc.variables

    def __getitem__(self, item):
        return self.nc.variables['PRCP_NC'][item] + \
               self.nc.variables['PRCP_C'][item]


class THETA(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['T'])
        self.units = 'K'
        self.description = 'Potential temperature'

    @staticmethod
    def can_do(nc):
        return 'T' in nc.variables

    def __getitem__(self, item):
        return self.nc.variables['T'][item] + 300.


class TK(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['T'])
        self.units = 'K'
        self.description = 'Temperature'

    @staticmethod
    def can_do(nc):
        return np.all([n in nc.variables for n in ['T', 'P', 'PB']])

    def __getitem__(self, item):
        p1000mb = 100000.
        r_d = 287.04
        cp = 7 * r_d / 2.
        t = self.nc.variables['T'][item] + 300.
        p = self.nc.variables['P'][item] + self.nc.variables['PB'][item]
        return (p/p1000mb)**(r_d/cp) * t


class PRESSURE(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['P'])
        self.units = 'hPa'
        self.description = 'Full model pressure'

    @staticmethod
    def can_do(nc):
        return np.all([n in nc.variables for n in ['P', 'PB']])

    def __getitem__(self, item):
        return self.nc.variables['P'][item] + self.nc.variables['PB'][item]


class GEOPOTENTIAL(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['PH'])
        self.units = 'm2 s-2'
        self.description = 'Full model geopotential'

    @staticmethod
    def can_do(nc):
        return np.all([n in nc.variables for n in ['PH', 'PHB']])

    def __getitem__(self, item):
        return self.nc.variables['PH'][item] + self.nc.variables['PHB'][item]


class SLP(FakeVariable):
    def __init__(self, nc):
        FakeVariable.__init__(self, nc)
        self._copy_attrs_from(nc.variables['T2'])
        self.units = 'hPa'
        self.description = 'Sea level pressure'
        dims = list(nc.variables['T'].dimensions)
        self.ds = np.nonzero(['bottom_top' in d for d in dims])[0][0]
        self._ds_shape = nc.variables['T'].shape[self.ds]

    @staticmethod
    def can_do(nc):
        # t2 is for attrs (not elegant)
        need = ['T', 'P', 'PB', 'QVAPOR', 'PH', 'PHB', 'T2']
        return np.all([n in nc.variables for n in need])

    def __getitem__(self, item):

        # take care of ellipsis and other strange indexes
        item = list(indexing.expanded_indexer(item, len(self.dimensions)))
        # we need the empty dims for _ncl_slp() to work
        squeezax = []
        for i, c in enumerate(item):
            if np.isscalar(c) and not isinstance(c, slice):
                item[i] = slice(c, c+1)
                squeezax.append(i)
        # add a slice in the 4th dim
        item.insert(self.ds, slice(0, self._ds_shape+1))
        item = tuple(item)

        # get data
        vars = self.nc.variables
        tk = vars['TK'][item]
        p = vars['P'][item] + vars['PB'][item]
        q = vars['QVAPOR'][item]
        z = (vars['PH'][item] + vars['PHB'][item]) / 9.81
        return np.squeeze(_ncl_slp(z, tk, p, q), axis=squeezax)

# Diagnostic variable classes in a list
var_classes = [cls.__name__ for cls in vars()['FakeVariable'].__subclasses__()]
var_classes.extend([cls.__name__ for cls in
                    vars()['AccumulatedVariable'].__subclasses__()])
var_classes.remove('AccumulatedVariable')


def _ncl_slp(z, t, p, q):
    """Computes the SLP out of the WRF variables.

    This code has been directly translated from the NCL fortran routine found
    in NCLS's wrf_user.f, therefore I reproduce their licence agreement below.

    Parameters
    ----------
    Z: geopotential height
    T: temp
    P: pressure
    Q: specific humidity

    NCL Licence
    -----------

    Copyright (C) 2015 University Corporation for Atmospheric Research
    The use of this software is governed by a License Agreement.
    See http://www.ncl.ucar.edu/ for more details.

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are
    met:

    Neither the names of NCAR's Computational and Information Systems
    Laboratory, the University Corporation for Atmospheric Research, nor
    the names of its contributors may be used to endorse or promote
    products derived from this Software without specific prior written
    permission.

    Redistributions of source code must retain the above copyright
    notice, this list of conditions, and the disclaimer below.

    Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions, and the disclaimer below in the
    documentation and/or other materials provided with the distribution.

    THIS SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
    EXPRESS OR IMPLIED, INCLUDING, BUT NOT LIMITED TO THE WARRANTIES OF
    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
    NONINFRINGEMENT. IN NO EVENT SHALL THE CONTRIBUTORS OR COPYRIGHT
    HOLDERS BE LIABLE FOR ANY CLAIM, INDIRECT, INCIDENTAL, SPECIAL,
    EXEMPLARY, OR CONSEQUENTIAL DAMAGES OR OTHER LIABILITY, WHETHER IN AN
    ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
    CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE
    SOFTWARE.
    """

    ndims = len(z.shape)
    if ndims == 4:
        out = []
        for _z, _t, _p, _q in zip(z, t, p, q):
            out.append(np.expand_dims(_ncl_slp(_z, _t, _p, _q), 0))
        return np.concatenate(out, axis=0)

    nx = z.shape[-1]
    ny = z.shape[-2]
    nz = z.shape[-3]

    r = 287.04
    g = 9.81
    gamma = 0.0065
    tc = 273.16 + 17.5
    pconst = 10000.

    #  Find least zeta level that is pconst Pa above the surface.  We
    # later use this level to extrapolate a surface pressure and
    # temperature, which is supposed to reduce the effect of the diurnal
    # heating cycle in the pressure field.

    p0 = p[0, ...]

    level = np.zeros((ny, nx), dtype=np.int) - 1
    for k in np.arange(nz):
        pok = np.nonzero((p[k, ...] < (p0 - pconst)) & (level == -1))
        level[pok] = k

    if np.any(level == -1):
        raise RuntimeError('Error_in_finding_100_hPa_up')  # pragma: no cover

    klo = (level-1).clip(0, nz-1)
    khi = (klo+1).clip(0, nz-1)

    if np.any((klo - khi) == 0):
        raise RuntimeError('Trapping levels are weird.')  # pragma: no cover

    x, y = np.meshgrid(np.arange(nx, dtype=np.int),
                       np.arange(ny, dtype=np.int))

    plo = p[klo, y, x]
    phi = p[khi, y, x]

    zlo = z[klo, y, x]
    zhi = z[khi, y, x]

    tlo = t[klo, y, x]
    thi = t[khi, y, x]

    qlo = q[klo, y, x]
    qhi = q[khi, y, x]

    tlo *= (1. + 0.608 * qlo)
    thi *= (1. + 0.608 * qhi)

    p_at_pconst = p0 - pconst
    t_at_pconst = thi - (thi-tlo) * np.log(p_at_pconst/phi) * np.log(plo/phi)
    z_at_pconst = zhi - (zhi-zlo) * np.log(p_at_pconst/phi) * np.log(plo/phi)
    t_surf = t_at_pconst * ((p0/p_at_pconst)**(gamma*r/g))
    t_sea_level = t_at_pconst + gamma * z_at_pconst

    # If we follow a traditional computation, there is a correction to the
    # sea level temperature if both the surface and sea level
    # temperatures are *too* hot.
    l1 = t_sea_level < tc
    l2 = t_surf <= tc
    l3 = ~l1
    t_sea_level = tc - 0.005 * (t_surf-tc)**2
    pok = np.nonzero(l2 & l3)
    t_sea_level[pok] = tc

    # The grand finale
    z_half_lowest = z[0, ...]

    # Convert to hPa in this step
    return 0.01 * (p0 * np.exp((2.*g*z_half_lowest)/(r*(t_sea_level+t_surf))))


def geogrid_simulator(fpath, do_maps=True):
    """Emulates geogrid.exe, which is useful when defining new WRF domains.

    Parameters
    ----------
    fpath: str
       path to a namelist.wps file
    do_maps: bool
       if you want the simulator to return you maps of the grids as well

    Returns
    -------
    (grids, maps) with:
        - grids: a list of Grids corresponding to the domains
          defined in the namelist
        - maps: a list of maps corresponding to the grids (if do_maps==True)
    """

    with open(fpath) as f:
        lines = f.readlines()

    pargs = dict()
    for l in lines:
        s = l.split('=')
        if len(s) < 2:
            continue
        s0 = s[0].strip().upper()
        s1 = list(filter(None, s[1].strip().replace('\n', '').split(',')))

        if s0 == 'PARENT_ID':
            parent_id = [int(s) for s in s1]
        if s0 == 'PARENT_GRID_RATIO':
            parent_ratio = [int(s) for s in s1]
        if s0 == 'I_PARENT_START':
            i_parent_start = [int(s) for s in s1]
        if s0 == 'J_PARENT_START':
            j_parent_start = [int(s) for s in s1]
        if s0 == 'E_WE':
            e_we = [int(s) for s in s1]
        if s0 == 'E_SN':
            e_sn = [int(s) for s in s1]
        if s0 == 'DX':
            dx = float(s1[0])
        if s0 == 'DY':
            dy = float(s1[0])
        if s0 == 'MAP_PROJ':
            map_proj = s1[0].replace("'", '').strip().upper()
        if s0 == 'REF_LAT':
            pargs['lat_0'] = float(s1[0])
        if s0 == 'REF_LON':
            pargs['lon_0'] = float(s1[0])
        if s0 == 'TRUELAT1':
            pargs['lat_1'] = float(s1[0])
        if s0 == 'TRUELAT2':
            pargs['lat_2'] = float(s1[0])

    # define projection
    if map_proj == 'LAMBERT':
        pwrf = '+proj=lcc +lat_1={lat_1} +lat_2={lat_2} ' \
             '+lat_0={lat_0} +lon_0={lon_0} ' \
             '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
        pwrf = pwrf.format(**pargs)
    elif map_proj == 'MERCATOR':
        pwrf = '+proj=merc +lat_ts={lat_1} ' \
             '+lon_0={lon_0} ' \
             '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
        pwrf = pwrf.format(**pargs)
    else:
        raise NotImplementedError('WRF proj not implemented yet: '
                                  '{}'.format(map_proj))
    pwrf = gis.check_crs(pwrf)

    # get easting and northings from dom center (probably unnecessary here)
    e, n = pyproj.transform(wgs84, pwrf, pargs['lon_0'], pargs['lat_0'])

    # LL corner
    nx, ny = e_we[0]-1, e_sn[0]-1
    x0 = -(nx-1) / 2. * dx + e  # -2 because of staggered grid
    y0 = -(ny-1) / 2. * dy + n

    # parent grid
    grid = gis.Grid(nxny=(nx, ny), ll_corner=(x0, y0), dxdy=(dx, dy),
                    proj=pwrf)

    # child grids
    out = [grid]
    for ips, jps, pid, ratio, we, sn in zip(i_parent_start, j_parent_start,
                                            parent_id, parent_ratio,
                                            e_we, e_sn):
        if ips == 1:
            continue
        ips -= 1
        jps -= 1
        we -= 1
        sn -= 1
        nx = we / ratio
        ny = sn / ratio
        if nx != (we / ratio):
            raise RuntimeError('e_we and ratios are incompatible: '
                               '(e_we - 1) / ratio must be integer!')
        if ny != (sn / ratio):
            raise RuntimeError('e_sn and ratios are incompatible: '
                               '(e_sn - 1) / ratio must be integer!')

        prevgrid = out[pid - 1]
        xx, yy = prevgrid.corner_grid.x_coord, prevgrid.corner_grid.y_coord
        dx = prevgrid.dx / ratio
        dy = prevgrid.dy / ratio
        grid = gis.Grid(nxny=(we, sn),
                        ll_corner=(xx[ips], yy[jps]),
                        dxdy=(dx, dy),
                        pixel_ref='corner',
                        proj=pwrf)
        out.append(grid.center_grid)

    maps = None
    if do_maps:
        from salem import Map
        import shapely.geometry as shpg

        maps = []
        for i, g in enumerate(out):
            m = Map(g)

            for j in range(i+1, len(out)):
                cg = out[j]
                left, right, bottom, top = cg.extent

                s = np.array([(left, bottom), (right, bottom),
                              (right, top), (left, top)])
                l1 = shpg.LinearRing(s)
                m.set_geometry(l1, crs=cg.proj, linewidth=(len(out)-j))

            maps.append(m)

    return out, maps
