#!/usr/bin/env python
# coding: utf-8

from datetime import datetime, timedelta

from xarray import DataArray, Dataset

import osr, ogr
import numpy as np
import scipy
import scipy.ndimage as ndi
import cartopy.crs as ccrs
import shapely.geometry

from smartcarb import config as cfg
from smartcarb import io


def rms(x):
    return np.sqrt(np.mean(x**2))

def source_in_swath(lon, lat, lon_o, lat_o, radius):
    s = get_circle(lon_o, lat_o, radius)
    points = np.transpose([lon.flatten(), lat.flatten()])
    m = shapely.geometry.MultiPoint(points)
    return m.intersects(s)



def list_swaths_at_source(filenames, origin, cloud_thr=1.0):

    lon_o, lat_o, radius = get_source_origin(origin)

    swath_fnames = []

    for filename in filenames:
        lon, lat, clct = io.read_fields(filename, ['longitude', 'latitude',
                                                      'CLCT'])
        mask = clct <= cloud_thr
        if source_in_swath(lon[mask], lat[mask], lon_o, lat_o, radius):

            swath_fnames.append(filename)

    return swath_fnames


def int_lon_eq_by_nsat(nsat, start=0.0):
    """
    Create integer list of equator starting positions for a constellation of
    `nsat` equally spaced satellites.
    """
    dlon = 24.146341463414636
    lon_eqs = (start + np.linspace(0.0, dlon, nsat+1)[:-1]) % dlon
    return lon_eqs.round().astype(int)


def lon_eq_by_nsat(nsat):
    DLON_ORBIT = 24.146341463414636
    lon_eqs = np.linspace(0, DLON_ORBIT, nsat+1)[:-1]
    lon_eqs = np.array([int(100 * round(v, 2)) for v in lon_eqs])
    return lon_eqs



def filter_dataset(ts, nsat=6, lon_eqs=None):

    if nsat is None:
        return ts

    if lon_eqs is None:
        lon_eqs = lon_eq_by_nsat(nsat)

    sats = ts.lon_eq == lon_eqs[0]

    for lon_eq in lon_eqs[1:]:
        sats |= np.abs(ts.lon_eq - lon_eq) < 1

    return ts[sats]



def interp_xy(xp,yp,x,y,c):
    """
    Get data value at x,y coords using linear interpolation.
    """
    valids = np.isfinite(c)
    x = x[valids]
    y = y[valids]
    c = c[valids]

    points = np.transpose([x, y])
    interp = scipy.interpolate.LinearNDInterpolator(points, c)

    return interp(xp,yp)



class DataField:
    def __init__(self, x, y, c, origin, crs='wgs84', pollon=cfg.pollon,
                 pollat=cfg.pollat, epsg=cfg.EPSG):
        """
        data field: c = f(x,y)

        transform: 'wgs84', 'rot' or 'xy'
        """
        self.data = c
        if hasattr(c, 'mask'):
            self.data[c.mask] = np.nan

        # coordinate system parameters
        self.origin = origin
        self.epsg = epsg
        self.pollon = pollon
        self.pollat = pollat

        # transform x,y to WGS 84
        if crs == 'wgs84':
            self.lon, self.lat = x, y
        elif crs == 'rot':
            self.lon, self.lat = rotpole2wgs(x, y, self.pollon, self.pollat)
        elif crs == 'xy':
            self.lon, self.lat = xy2wgs(x, y, origin, epsg)
        else:
            raise ValueError('Unknown crs "%s".' % crs)

        self.update()


    def update(self):
        """ Update x,y and and rlon, rlat."""
        print('update only if x and y is needed')

        if self.epsg is None:
            self.x, self.y = None, None
        else:
            self.x, self.y = xy2wgs(self.lon, self.lat, self.origin, self.epsg,
                                    inverse=True)
        if self.pollon is None:
            self.rlon = None
            self.rlat = None
        else:
            self.rlon, self.rlat = rotpole2wgs(self.lon, self.lat, self.pollon,
                                               self.pollat, inverse=True)

    def interp_xy(self, x,y):
        """
        Get data value at x,y coords using linear interpolation.
        """
        points = np.transpose([self.x.flatten(), self.y.flatten()])
        interp = scipy.interpolate.LinearNDInterpolator(points, self.flatten())

        return interp(x,y)


    def fill_gaps(self, gaps=None):
        """
        Fill gaps in data array. Gaps are all non finite values in data array
        plus gaps given by gaps array.
        """
        valids = np.isfinite(self)

        if gaps is not None:
            valids[gaps] = False

        points = np.transpose([self.x[valids], self.y[valids]])
        interp = scipy.interpolate.LinearNDInterpolator(points,
                                                        self[valids])
        return interp(self.x, self.y)



def get_source_origin(name):
    if name == 'Berlin':
        lon_o, lat_o = cfg.CITIES[name]
        radius = 15.0 # km
    else:
        lon_o, lat_o = cfg.POINT_SOURCES[name]
        radius = 5.0 # km

    return lon_o, lat_o, radius


def find_all_roots(x,y):
    """
    Find all roots for y = f(x) assuming
    line between x[i] and x[i+1].
    """
    x = np.asarray(x)
    y = np.asarray(y)
    roots = []

    for i in range(x.size-1):
        m = (y[i+1] - y[i]) / (x[i+1] - x[i])
        b = y[i] - m * x[i]
        x0 = -b / m

        if x[i] <= x0 < x[i+1]:
            roots.append(x0)

    return np.array(roots)


def get_circle(lon_o, lat_o, radius=5.0):
    """
    Return circle around origin with radius (in km).
    """
    radius = radius / 100.0 # km -> degree (roughly)
    p = shapely.geometry.Point(lon_o, lat_o)

    if radius > 0:
        p = p.buffer(radius)

    return p




def compute_centroid(x, y, w):
    """\
    Compute center of plume.
    """
    x0 = np.sum(w * x) / np.sum(w)
    y0 = np.sum(w * y) / np.sum(w)

    return x0, y0



def anglediff(alpha, beta, is_degree=True):
    if is_degree:
        alpha = np.deg2rad(alpha)
        beta = np.deg2rad(beta)

    delta = np.arctan2(
        np.sin(alpha-beta),
        np.cos(alpha-beta)
    )
    return np.rad2deg(delta) if is_degree else delta


def nearest_index(lon_o, lat_o, lon, lat):
    index = np.argmin(np.abs(lon - lon_o) + np.abs(lat - lat_o))
    i,j = np.unravel_index(index, lon.shape)

    return i,j


def wind_direction(u, v):
    return (270.0 - np.rad2deg(np.arctan2(v,u))) % 360


def iter_times(start, end, step=timedelta(hours=1), exclude_end=False):
    current = start

    if exclude_end:
        while current < end:
            yield current
            current = current + step
    else:
        while current <= end:
            yield current
            current = current + step



def create_rotpole_reference(pollon, pollat):
    """\
    Return spatial reference for rotated pole coordinates.

    pollon: longitude of north pole
    pollat: latitude of north pole
    """
    ref = osr.SpatialReference()
    ref.ImportFromEPSG(4326)     # WSG-84

    if (pollon is not None) and (pollat is not None):
        o_lon_p = 0.0
        o_lat_p = pollat
        lon_0 = 180 + pollon

        string = (
            'PROJCS["rotated_pole projection",EXTENSION["PROJ4","+proj=ob_tran +o_proj=longlat'
            ' +o_lon_p=%.2f +o_lat_p=%.2f +lon_0=%.2f +a=1 +to_meter=0.0174532925199"]]'
        ) % (o_lon_p, o_lat_p, lon_0)

        ref.ImportFromWkt(string)    # rotpole transformation

    return ref




def xy2wgs(x, y, origin, epsg=cfg.EPSG, inverse=False, use_xarray=True):
    """
    xy (meters from origin) to WGS84.
    """
    if use_xarray:
        dims = x.dims

    lon_o, lat_o, _ = get_source_origin(origin)
    x_o, y_o = wgs2epsg(lon_o, lat_o, epsg)

    if inverse:
        attrs = {'origin': (lon_o, lat_o), 'epsg': epsg, 'origin_name': origin}
        x,y = wgs2epsg(x.data,y.data, epsg, inverse=False)
        x = x - x_o
        y = y - y_o

        if use_xarray:
            x = DataArray(x, name='x', dims=dims, attrs=attrs)
            y = DataArray(y, name='y', dims=dims, attrs=attrs)

        return x,y

    else:
        x = x + x_o # in m
        y = y + y_o # in m

        lon, lat = wgs2epsg(x,y, epsg, inverse=True)

        if use_xarray:
            lon = DataArray(lon, name='lon', dims=dims)
            lat = DataArray(lat, name='lat', dims=dims)

        return lon, lat


def wgs2epsg(lon, lat, epsg=cfg.EPSG, inverse=False):
    """
    Transforms lon/lat to EPSG.
    """
    if inverse:
        out = ccrs.PlateCarree()
        in_ = ccrs.epsg(epsg)
    else:
        out = ccrs.epsg(epsg)
        in_ = ccrs.PlateCarree()

    if np.ndim(lon) == 0:
        res = out.transform_point(lon, lat, in_)
        return res[0], res[1]
    elif np.ndim(lon) in [1,2]:
        res = out.transform_points(in_, lon, lat)
        return res[...,0], res[...,1]
    else:
        shape = lon.shape
        res = out.transform_points(in_, lon.flatten(), lat.flatten())
        return res[:,0].reshape(shape), res[:,1].reshape(shape)


def rotpole2wgs(rlon, rlat, pollon=cfg.pollon, pollat=cfg.pollat,
                inverse=False):
    """
    Transform rotated pole to WGS84.
    """
    c_in = ccrs.RotatedPole(pollon, pollat)
    c_out  = ccrs.PlateCarree()

    if inverse:
        c_in, c_out = c_out, c_in

    if np.ndim(rlon) == 0:
        res = c_out.transform_point(rlon, rlat, c_in)
        return res[0], res[1]
    elif np.ndim(rlon) in [1,2]:
        res = c_out.transform_points(c_in, rlon, rlat)
        return res[...,0], res[...,1]
    else:
        shape = rlon.shape
        res = c_out.transform_points(c_in, rlon.flatten(), rlat.flatten())
        return res[:,0].reshape(shape), res[:,1].rehape(shape)



def spatial_binning(a, n=2):
    if np.ndim(a) == 2:
        kernel = np.ones((n,n)) / n**2
        return ndi.convolve(a, kernel)[::n,::n]
    else:
        kernel = np.ones((1,n,n)) / n**2
        return ndi.convolve(a, kernel)[:,::n,::n]


