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

from __future__ import print_function

from scipy.interpolate import interp1d

import itertools
import sys

from matplotlib.colors import LinearSegmentedColormap, LogNorm
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import netCDF4
import shapely
import rasterio.features

import cartopy
import cartopy.crs as ccrs
import cartopy.feature as cfeature

import c2sm
import orbitsim as sim

#from . import orbit
from smartcarb import misc
from smartcarb import orbit
from smartcarb import config as cfg



BERLIN = 13.407696, 52.516984

POWER_STATION = cfg.POINT_SOURCES
#{
#    'Jänschwalde':    (14.457934, 51.836579),
#    'Schwarze Pumpe': (14+21/60+12/3600, 51+32/60+10/3600),
#    'Boxberg':        (14+33/60+53/3600, 51+24/60+58/3600),
#    'Lippendorf':     (12+22/60+40/3600, 51+11/60+ 7/3600),
#    'Turów':          (14.913, 50.948),
#    'Patnów':         (18.236118, 52.301086)
#}

# with 0.05 border on Smartcarb domain
XLIM = (-1.0554999649524688, 5.2455001413822178)
YLIM = (2.7944999933242798, 8.1954997777938843)



def add_scatter(ax, lon, lat, slice=Ellipsis, marker='o', color='k', ms=1.0,
                label=None):

    x = np.array(lon)[slice]
    y = np.array(lat)[slice]

    return ax.plot(x, y, marker=marker, color=color, ms=ms, ls='',
                   transform=ccrs.PlateCarree(), label=label)




def add_curve(ax, x, y, origin, epsg=cfg.EPSG, label=None,
              color=None, ls='-'):
    """\
    Add a curve to map.
    """
    if epsg is None:
        lines = ax.plot(x, y, ls=ls, color=color, label=label)
    else:
        lon, lat = misc.xy2wgs(x, y, origin, epsg=epsg, use_xarray=False)
        lines = ax.plot(lon, lat, ls=ls, color=color, transform=ccrs.PlateCarree(),
                        label=label)

    if label is not None:
        plt.legend()


def add_city_labels(ax, color='black', size='medium', suffixes=None, names=None,
                    va='flex', ha='left', fontdict=None):

    if names is None:
        names = ['Berlin'] + sorted(POWER_STATION.keys())

    if 'Berlin' in names:
        if isinstance(ha, str):
            align = ha
        else:
            align = ha[0]

        names.remove('Berlin')
        ax.plot(BERLIN[0], BERLIN[1], '.', color=color,
                transform=ccrs.PlateCarree(), label='_none_')
        ax.text(BERLIN[0], BERLIN[1], ' Berlin ', va='center', ha=align, fontsize=size,
                transform=ccrs.PlateCarree(), color=color, fontdict=fontdict)

    for i, key in enumerate(names, 1):
        lon, lat = POWER_STATION[key]
        ax.plot(lon, lat, '.', color=color, transform=ccrs.PlateCarree(),
                label='_none_')

        if isinstance(ha, str):
            align = ha
        else:
            align = ha[i]

        if key in ['Patnów'] and va == 'flex':#, 'Schwarze Pumpe']:#, 'Boxberg']:
            ax.text(lon, lat, u'%s  ' % key, va='center', ha='right', fontsize=size,
                    transform=ccrs.PlateCarree(), color=color, clip_on=True,
                    fontdict=fontdict)
        else:
            if align == 'left':
                key = u' %s' % key
            elif align == 'right':
                key = u'%s ' % key

            ax.text(lon, lat, key, va='center', ha=align, fontsize=size,
                    transform=ccrs.PlateCarree(), color=color, clip_on=True,
                    fontdict=fontdict)


def draw_smartcarb_domain(ax, color='black'):
    projection = ccrs.RotatedPole(
        pole_latitude=cfg.pollat,
        pole_longitude=cfg.pollon
    )
    lon = [cfg.startlon, cfg.startlon, cfg.stoplon, cfg.stoplon, cfg.startlon]
    lat = [cfg.startlat, cfg.stoplat, cfg.stoplat, cfg.startlat, cfg.startlat]

    ax.plot(lon, lat, '-', color=color, transform=projection)



def create_orbits(scenario_filename, start=0, length=1, lon_eq=None, nrow=None):
    """\
    Create list of orbit polgons for instrument scenario. `nrow` is the
    number of across-track pixels. A small number (nrow = 2) will be fast
    but might fail for wide swaths because the viewing angles for the corner
    points might be larger than 90 degree.
    """
    # init orbit simulator
    s = sim.Orbit_sim.from_file(scenario_filename)
    s.tsamp = 10.0
    s.nrow = s.nrow if nrow is None else nrow

    if lon_eq is not None:
        s.longitude_eq = np.deg2rad(lon_eq)

    orbits = []
    days = []
    for i in range(start, start+length):
        s.phase_start = i + 0.29 # 0.33#0.29
        s.phase_end = i + 0.71   # 0.37#0.71
        s.calculate_revolution()

        orbit = create_orbit_polygon(s)
        orbits.append(orbit)
        days.append(s.utc.mean(0)[2])

    return days, orbits




def create_orbit_polygon(s):
    """\
    Create a polygon from the boundary of orbit (using center points).
    """
    lon, lat, lon_center, lat_center = orbit.reshape_orbit(s)
    nrows, ncols = lon.shape

    x = np.concatenate([lon[0,:], lon[:,-1], lon[-1,::-1], lon[::-1,0]])[::-1]
    y = np.concatenate([lat[0,:], lat[:,-1], lat[-1,::-1], lat[::-1,0]])[::-1]
    coords = np.transpose([x,y])
    orbit_poly = shapely.geometry.Polygon(coords)

    if not orbit_poly.is_valid or lon[-1,ncols//2] > lon[0,ncols//2]:
        lon[lon < 0] = 360.0 + lon[lon < 0]
        #lon[lon > 180] = lon[lon > 180] - 360.0
        x = np.concatenate([lon[0,:], lon[:,-1], lon[-1,::-1], lon[::-1,0]])[::-1]
        y = np.concatenate([lat[0,:], lat[:,-1], lat[-1,::-1], lat[::-1,0]])[::-1]
        coords = np.transpose([x,y])
        orbit_poly = shapely.geometry.Polygon(coords)

        if not orbit_poly.is_valid:
            print('Orbit not valid')

    return orbit_poly


def draw_polygon(polygon, ax, edgecolor='none', fillcolor='silver', alpha=0.5,
        transform=None, zorder=-1):
    """\
    Draw an orbit polygon on matplotlib axes.
    """
    if transform is None:
        transform = ccrs.PlateCarree()

    feature = cartopy.feature.ShapelyFeature([polygon], transform)
    ax.add_feature(feature, color=fillcolor, alpha=alpha, edgecolor=edgecolor, zorder=zorder)
    return feature


def rasterize_polygon(polygon):
    shape = (int(180/resolution), int(360/resolution))
    transform = rasterio.Affine(resolution,0,-180,0,resolution,-90)
    rasterio.features.rasterize([polygon], shape, transform=transform)


def rasterize(orbits, resolution=0.01, lon_min=-180, lon_max=180, lat_min=-90,
             lat_max=90):
    """\
    Rasterize orbits on lon/lat grid and returns array with number of overlaps.
    """
    # lon = [-180,180]
    shape = (int((lat_max-lat_min)/resolution), int((lon_max-lon_min)/resolution))
    transform = rasterio.Affine(resolution,0,lon_min,0,resolution,lat_min)
    img = np.array([rasterio.features.rasterize([orbit], shape, transform=transform)
         for orbit in orbits])

    # lon = [180,360]
    transform = rasterio.Affine(resolution,0,lon_min+360,0,resolution,lat_min)
    img2 = np.array([rasterio.features.rasterize([orbit], shape, transform=transform)
        for orbit in orbits])

    return (img+img2).sum(0)


def calculate_pixel(scenario_filename, orbit_number):
    s = sim.Orbit_sim.from_file(scenario_filename)
    s.phase_start = orbit_number + 0.349
    s.phase_end = orbit_number + 0.354
    s.calculate_revolution()
    return orbit.reshape_orbit(s)


def make_orbit_map(scenario_filename, domain, orbit=0, length=14, nrow=11,
        add_labels=False, add_domain=False, add_pixels=False, do_overlap=False,
        vmin=None, vmax=None, size='large', color='white', cmap='viridis',
        orbits=None):

    if orbits is None:
        days, orbits = create_orbits(scenario_filename, orbit, length, nrow=nrow)

    domain = cfg.DOMAINS[domain]
    admin_level = 1 if domain.name in ['smartcarb', 'berlin'] else 0
    fig, ax, cax = c2sm.vis.create_map(domain, add_colorbar=do_overlap,
                                       admin_level=admin_level)

    if add_labels:
        if domain.name == 'smartcarb':
            add_city_labels(ax, color=color)
        else:
            add_city_labels(ax, color=color, names=['Berlin', 'Jänschwalde', 'Patnów'], va='flex')

    if add_domain:
        draw_smartcarb_domain(ax, color=color)

    if add_pixels:
        s = sim.Orbit_sim.from_file(scenario_filename)
        s.phase_start = orbit + 0.351
        s.phase_end = orbit + 0.354
        #s.nrow = 2
        s.calculate_revolution()
        lon, lat, lon_center, lat_center = sim.reshape_orbit(s)
        print(lon.shape, lon_center.shape)

        m,n = lon.shape

        for i in range(m):
            ax.plot(lon_center[i,:,0], lat_center[i,:,0], 'k-',
                    transform=ccrs.PlateCarree(),
                    linewidth=0.5, color='blue'
            )
        for j in range(n):
            ax.plot(lon_center[:,j,0], lat_center[:,j,0], 'k-',
                    transform=ccrs.PlateCarree(),
                    linewidth=0.5, color='blue'
            )

    if do_overlap:
        if domain.name == 'smartcarb':
            raster = rasterize(orbits, resolution=0.01, lon_min=7, lon_max=20,
                               lat_min=49, lat_max=56)
            extent = [7,20,49,56]

        else:
            raster = rasterize(orbits, resolution=0.1, lon_min=-180, lon_max=180,
                                   lat_min=-90, lat_max=90)
            extent = [-180,180,-90,90]

        vmin = np.min(raster) if vmin is None else vmin
        vmax = np.max(raster) if vmax is None else vmax

        cmap = plt.cm.get_cmap(cmap, vmax-vmin+1)
        c = ax.imshow(raster, vmin=vmin-0.5, vmax=vmax+0.5, cmap=cmap,
                      extent=extent, alpha=0.9,
                      transform=ccrs.PlateCarree())
        cb = plt.colorbar(c, cax=cax)
        cb.set_label('number of overlapping swaths', fontsize=size)
        cb.set_ticks(np.arange(vmin, vmax+1))
    else:
        for orbit in orbits:
            pass#draw_orbit(orbit, ax)


def make_cosmo_map(rlon, rlat, field, vmin=None, vmax=None, cmap='viridis',
        border=0.00, norm='linear', label='', domain='smartcarb', fig=None,
        clct=None, xlim=None, ylim=None, edgecolor='black',
        add_colorbar=True, fig_size=None):

    domain = cfg.DOMAINS[domain]
    return c2sm.vis.make_field_map(rlon, rlat, field, domain=domain, vmin=vmin,
                                   vmax=vmax, cmap=cmap, border=border, norm=norm,
                                   label=label, fig=fig, fig_size=fig_size, alpha=clct,
                                   xlim=xlim, ylim=ylim, edgecolor=edgecolor,
                                   add_colorbar=add_colorbar)


def load_orbit_data(filename, tracer):
    with netCDF4.Dataset(filename) as nc:
        lon = nc.variables['longitude_corners'][:]
        lat = nc.variables['latitude_corners'][:]

        if tracer == 'XCO2':
            field1 = nc.variables['XCO2_TOT'][:]
            field2 = nc.variables['XCO2_GPP'][:]
            field3 = nc.variables['XCO2_RA'][:]
            field4 = nc.variables['XCO2_BG'][:]

            field = field1 - field2 + field3 + field4

        elif tracer == 'NO2':
            field = nc.variables['NO2_TOT_column'][:] + nc.variables['NO2_BG_column'][:]

        elif tracer == 'CO':
            field = nc.variables['CO_TOT_column'][:] + nc.variables['CO_BG_column'][:]

        else:
            field = nc.variables[tracer][:]

    return lon, lat, field


def make_level2_map(lon, lat, values, fig=None, vmin=None, vmax=None,
        domain='smartcarb', label='', alpha=1.0, cmap='viridis',
        clct=None, xlim=XLIM, ylim=YLIM, edgecolor='black',
        is_discrete=False, cb_labels=None, cb_labelsize='small',
        truncate_cmap=False, do_zoom=False, zoom_area=None, fig_size=6.0,
                    dright=0.0, bg_color='silver'):
    """\
    lon, lat :: longitude and latitude of pixel corners
    values   :: 2d field of level-2 data

    is_discrete :: bool (default False)
        if True uses colormap with discrete levels with one colour per values
        in fields
    """
    domain = cfg.DOMAINS[domain]

    return c2sm.vis.make_level2_map(lon, lat, values, domain=domain, fig=fig,
                                    vmin=vmin, vmax=vmax, label=label,
                                    alpha=alpha, do_zoom=do_zoom,
                                    zoom_area=zoom_area,
                                    cmap=cmap, clct=clct, xlim=xlim, ylim=ylim,
                                    edgecolor=edgecolor, is_discrete=is_discrete,
                                    cb_labels=cb_labels, cb_labelsize=cb_labelsize,
                                    truncate_cmap=truncate_cmap,
                                    fig_size=fig_size, dright=dright,
                                    bg_color=bg_color)





def calculate_xco2(xm, p, ps, h):

    p = p[:,::-1,:,:]
    h = p[:,::-1,:,:]
    xm = p[:,::-1,:,:]

    _, ke, je, ie = xm.shape

    xco2 = np.zeros((je,ie))

    for i,j in itertools.product(range(ie), range(je)):
        h_mid = 0.5 * (h[0,1:,j,i] + h[0,:-1,j,i])
        pp = np.interp(np.log(h[0,:,j,i]), np.log(h_mid), p[0,:,j,i])
        pp[0] = ps[0,j,i]

        mair = - np.diff(pp) / g
        xco2[j,i] = np.sum(mair * xm[0,j,i]) / np.sum(mair)

    return xco2





def main(field, rlon, rlat, rotpole=None, vmin=None, vmax=None, cmap=None,
        label=None, edgecolor='white', border=50, ax=None, img=None,
        title=None):
    """\
    Make map. `ax` and `img` can be already drawn map of which only
    the image is replaced with a new one (assuming same vmin, vmax etc.).
    """

    if label is None:
        label = 'no title'

    dlon = (rlon[1]-rlon[0])
    left = rlon[0] - dlon / 2.0
    right = rlon[-1] + dlon / 2.0

    dlat = (rlat[1]-rlat[0])
    bottom = rlat[0] - dlat / 2.0
    top = rlat[-1] + dlat / 2.0

    make_new = (ax is None and img is None)

    if make_new:

        if rotpole is None:
            rotated_pole = ccrs.RotatedPole(
                pole_latitude=cfg.pollat,
                pole_longitude=cfg.pollon
            )
        else:
            rotated_pole = ccrs.RotatedPole(
                pole_latitude=rotpole['grid_north_pole_latitude'],
                pole_longitude=rotpole['grid_north_pole_longitude']
            )

        fig = plt.figure(figsize=(5.8, 4))
        ax = plt.subplot(111, projection=rotated_pole)

        ax.set_xlim(left + border, right - border)
        ax.set_ylim(bottom + border, top - border)

        # add coastlines and countries
        ax.coastlines(resolution='10m', color=edgecolor, linewidth=0.5)
        lines = cfeature.NaturalEarthFeature(
            category='cultural',
            name='admin_0_boundary_lines_land',
            scale='10m',
        )
        ax.add_feature(lines, edgecolor=edgecolor, facecolor='none', linewidth=0.5)
        #lines = cfeature.NaturalEarthFeature(
        #    category='cultural',
        #    name='admin_1_states_provinces_lines',
        #    scale='10m',
        #)
        #ax.add_feature(lines, edgecolor=edgecolor, facecolor='none', linewidth=0.5)

    else:
        img.remove()

    # show image
    if title is not None:
        ax.set_title(title)

    img = ax.imshow(field,
        extent=(left, right, bottom, top),
        vmin=vmin, vmax=vmax, cmap=cmap
    )

    if make_new:
        cb = plt.colorbar(img, extend='both')
        cb.set_label(label)

        plt.tight_layout()
        plt.show()


    return ax, img


if __name__ == '__main__':


    make_orbit_map('config/Sentinel-7-XCO2.txt', 'berlin', 0, 1,
            add_labels=True, add_domain=True, do_overlap=False, nrow=None,
            color='black', add_pixels=True
    )
    plt.savefig('../figures/Sentinel-7-berlin.png')
    plt.savefig('../figures/Sentinel-7-berlin.pdf')
    plt.close()

    sys.exit(0)

    orbits = create_orbits('config/Sentinel-7-XCO2.txt', 0, 1, nrow=2)[1]
    orbits += create_orbits('config/Sentinel-7-NO2_4km.txt', 0, 1, nrow=2)[1]
    make_orbit_map(None, 'smartcarb', 0, 1,
            add_labels=True, add_domain=True, do_overlap=False, nrow=2,
            color='black', orbits=orbits
    )
    plt.savefig('../figures/Sentinel-7-orbit.png')
    plt.savefig('../figures/Sentinel-7-orbit.pdf')
    plt.close()

    make_orbit_map('config/Sentinel-7-XCO2.txt', 'europe', 0, 164,
            add_labels=True, add_domain=True, do_overlap=True, nrow=2,
            vmin=0, vmax=6
    )
    plt.savefig('../figures/Sentinel-7-XCO2-coverage.png')
    plt.savefig('../figures/Sentinel-7-XCO2-coverage.pdf')
    plt.close()


    make_orbit_map('config/Sentinel-7-NO2_4km.txt', 'europe', 0, 164,
            add_labels=True, add_domain=True, do_overlap=True, nrow=2,
            vmin=0, vmax=6
    )
    plt.savefig('../figures/Sentinel-7-NO2-coverage.png')
    plt.savefig('../figures/Sentinel-7-NO2-coverage.pdf')
    plt.close()















