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

from datetime import datetime, timedelta


from collections import defaultdict
from fnmatch import fnmatch
import logging
import os
import time

import cartopy.crs as ccrs
from netCDF4 import Dataset

import numpy as np
import shapely

from smartcarb import orbit


from . import config as cfg
from . import epf
from . import io
from . import misc


np.seterr(divide='ignore', invalid='ignore')



def get_var_names(var_patterns, var_names):
    """\
    Get variable names that match patterns.
    """
    names = []
    for pattern in var_patterns:
        names += [key for key in var_names if fnmatch(key, pattern)]

    return sorted(set(names))


def read_variables(filename, utc, var_names, attr_names, utc0=datetime(2015,1,1)):
    """\
    Read variables and attributes from netCDF file. If file has a
    'time' variable use nearest time.
    """
    t = (utc - utc0).total_seconds()

    fields = []
    attrs = []

    logging.debug('open %s' % filename)
    with Dataset(filename) as nc:
        var_names = get_var_names(var_names, nc.variables.keys())

        if 'time' in nc.variables:
            index = np.argmin(np.abs(nc.variables['time'][:] - t))
        else:
            index = Ellipsis

        for name in var_names:
            v = nc.variables[name]
            fields.append(v[index])
            attrs.append(
                dict((k, v.getncattr(k)) for k in attr_names)
            )

    return fields, attrs, var_names



def iter_pixel_area(lon, lat, x0=cfg.startlon, y0=cfg.startlat,
        dx=cfg.dlon, dy=cfg.dlat, ie=cfg.ie, je=cfg.je, grid_cells=None):
    """\
    Iterate over indices of grid cells and yield fraction of overlap
    of pixel and cell.
    """
    if grid_cells is None:
        grid_cells = dict()


    pixel = create_pixel(lon, lat)

    if not pixel.intersects(cfg.DOMAIN):
        raise StopIteration

    dx2, dy2 = 0.5 * dx, 0.5 * dy

    # polygon bounds
    minx, miny, maxx, maxy = pixel.bounds

    # round up and down
    minx = np.floor(minx / dx) * dx
    miny = np.floor(miny / dy) * dy
    maxx = np.ceil(maxx / dx) * dx + 0.5 * dx
    maxy = np.ceil(maxy / dy) * dy + 0.5 * dy

    # start point for pixel
    i0 = int(round((minx - x0) / dx))
    j0 = int(round((miny - y0) / dy))

    # iter over grid cells
    for i,x in enumerate(np.arange(minx,maxx,dx), i0):
        for j,y in enumerate(np.arange(miny,maxy,dy), j0):

            if 0 <= i < ie and 0 <= j < je:

                # cache grid cells for re-using
                if (i,j) in grid_cells:
                    grid_cell = grid_cells[i,j]
                else:
                    grid_cell = shapely.geometry.Polygon([
                        [x-dx2, y-dy2],
                        [x-dx2, y+dy2],
                        [x+dx2, y+dy2],
                        [x+dx2, y-dy2]
                    ])
                    grid_cells[i,j] = grid_cell

                # calculate fraction of intersection and yield
                if grid_cell.intersects(pixel):
                    weight = grid_cell.intersection(pixel).area / grid_cell.area
                    yield j, i, weight



def create_pixel(lon, lat, pollon=cfg.pollon, pollat=cfg.pollat):
    """\
    Create pixel polygon in rotpole coords from corner longitude
    and latitude.
    """
    inp = ccrs.PlateCarree()
    out = ccrs.RotatedPole(
        pole_latitude=pollat,
        pole_longitude=pollon
    )
    coords = out.transform_points(inp, lon, lat)
    pixel = shapely.geometry.Polygon(coords)

    return pixel


def pre_compute_indices(lon, lat):

    indices = dict()

    # cache for grid cell polygons
    grid_cells = dict()

    for n,m in np.ndindex(lon.shape[:2]):
        indices[n,m] = {
            'i': [],
            'j': [],
            'w': []
        }
        for j,i,weight in iter_pixel_area(lon[n,m], lat[n,m], grid_cells=grid_cells):
            indices[n,m]['i'].append(i)
            indices[n,m]['j'].append(j)
            indices[n,m]['w'].append(weight)

    return indices



def sum_satellite_pixels(fields, lon, lat, indices=None):
    """\
    Create satellite observations from orbit filename.

    fields   :: data fields as list of 2D/3D arrays
    lon, lat :: longitude and latitude of satellite pixel corners
                with shape: (nobs, npixel, ncorners)

    iterator :: pre-computed indices for re-using
    """
    n_fields = len(fields)
    shape = lon.shape[:2]

    values = np.zeros((n_fields, shape[0], shape[1]), dtype='f4')
    weights = np.zeros(shape, dtype='f4')

    for n,m in np.ndindex(shape):

        if indices is None:
            iterator = iter_pixel_area(lon[n,m], lat[n,m])
        else:
            iterator = zip(indices[n,m]['j'], indices[n,m]['i'],
                    indices[n,m]['w'])

        for j,i,weight in iterator:
            for k in range(n_fields):

                if np.ndim(fields[k]) == 3:
                    values[k,n,m] += fields[k][0,j,i] * weight
                else:
                    values[k,n,m] += fields[k][j,i] * weight

            weights[n,m] += weight

    for k in range(n_fields):
        values[k] /= weights

    values = np.ma.array(values, mask=np.isnan(values))

    return values



def read_orbit(orbit_filename):
    with Dataset(orbit_filename) as nc:
        time = nc.variables['time'][...]
        lon0 = nc.variables['longitude'][...]
        lat0 = nc.variables['latitude'][...]
        lonc = nc.variables['longitude_corners'][...]
        latc = nc.variables['latitude_corners'][...]
        sza =  nc.variables['solar_zenith_angle'][...]

    return time, lon0, lat0, lonc, latc, sza


def date2orbit(date, instrument=None):
    """\
    """
    cycle_duration, cycle_length = {
                       #  days, orbits in cycle
            'Sentinel_5': (29,  412),
            'Sentinel_7': (11,  164)
    }[instrument]

    days = (date - datetime(2015,1,1)).total_seconds() / 86400.0
    return round(days / cycle_duration * cycle_length)



def calculate_sza():
    # TODO
    pass


def any_pixels_in_domain(lon, lat, domain=cfg.DOMAIN, domain_projection=cfg.ROTPOLE):
    """\
    Check if any (lon, lat) intersects with model domain.
    """
    lon, lat = lon.flatten(), lat.flatten()
    points = domain_projection.transform_points(cfg.WGS84, lon, lat)
    points = shapely.geometry.MultiPoint(points)
    return points.intersects(cfg.DOMAIN)



def check_overpass(instrument, lon_eq, cycle_length=165):

    orbits = []

    for n in range(0,cycle_length):
        print(n)
        orbit_data = orbit.compute_orbit(instrument, n,
                long_eq=lon_eq, do_reshape=True, write_to_file=False
        )
        if any_pixels_in_domain(orbit_data.lon_center, orbit_data.lat_center):
            orbits.append(n)

    return orbits


def main(instrument, output_path, dont_overwrite=False,
        cosmo_path=None, cosmo_vars=None, cosmo_attrs=None,
        albedo_path=None, albedo_vars=None, albedo_attrs=None,
        overpass_time=None, compute_orbit=False,
        orbit_path=None, orbit_number=0, lon_eq=0.0,
        do_shift_orbit=False,
    ):
    """\
    Create Level-2 data.

    Arguments:

    - instrument :: str
       name of instrument scenario (used for filename and pick default arguments
       defined in config.py)

    - output_path :: str
        path data output

    - dont_overwrite :: bool (default: False)
        if True don't overwrite existing files

    - cosmo_path :: str
        Path to COSMO 2D file. If the path has datetime formats (e.g. %Y%m%d%H)
        `utc_time` will be used to format string.

    - cosmo_vars :: list of str (defaults in config.py)
        list of tracer names that will be taken from COSMO 2D file and
        averaged on instrument pixels. It is possible to use Unix filename
        pattern matching. E.g. "XCO2_*" matches all XCO2 variables.

    - cosmo_attrs :: list of str (defaults in config.py)
        names of COSMO attributes that will be copied from the netCDF variables
        in the COSMO file.

    - albedo_path :: str (default: None)
        path to albedo filename. If the path as datetime formats (e.g. %Y%m%d%H)
        `utc_time` will be used to format string.

    - albedo_vars :: str (defaults in config.py)
        list of albedo variables that will be taken from the albedo file. It is possible
        to use Unix filename pattern matching.

    - albedo_attrs :: list of str (defaults in config.py)
        names of attributes that will be copied from the netCDF variables
        in the albedo file.

    - overpass_time :: datetime (default: None)
        explicit overpass time that will overwrite the time
        from orbit file/computation

    - compute_orbit :: boolean (default: False)
        compute orbit online (requires SRON orbit simulator)

    - orbit_path :: str
        netCDF file with orbit data (e.g. lon/lat of pixels)
        OR file describing orbit parameters

        If file describes orbit parameter, `online_orbit` needs
        to be `True` and `instrument`, `orbit_number` and `lon_eq`
        needs to be provided.

    - do_shift_orbit :: boolean (default: False)
       shift orbit using orbit number and cycle duration of given
       instrument (not fully implemented)
    """
    if instrument is None:
        raise ValueError('Please give "instrument" name.')

    # compute or read orbit
    if compute_orbit:

        logging.info('Compute %s orbit #%d (lon_eq: %d)' % (instrument, orbit_number, lon_eq))
        orbit_data = orbit.compute_orbit(instrument, orbit_number,
                long_eq=lon_eq, do_reshape=True, write_to_file=False
        )
        times = orbit_data.times
        lon0 = orbit_data.lon_center
        lat0 = orbit_data.lat_center
        lonc = orbit_data.lon_corner
        latc = orbit_data.lat_corner
        sza = orbit_data.sza

    # shift pre-computed orbit
    elif do_shift_orbit:
        logging.info('Read orbit and shift.')
        logging.warning('Shifting orbits not fully implemented.')

        # get cycle and orbit number
        if instrument == 'Sentinel_5':
            cycle_duration = 29.0 # days
            cycle_number, number  = divmod(orbit_number, 412)
        else:
            cycle_duration = 11.0 # days
            cycle_number, number  = divmod(orbit_number, 164)

        # TODO: should make sure that orbit path has a number format string
        orbit_path = orbit_path.format(number=number)
        times, lon0, lat0, lonc, latc, sza = read_orbit(orbit_path)

        # TODO: update time and SZA
        times[:] = np.nan
        sza[:] = np.nan

        # shift equator passing longitude (e.g. 0 or 180) for two satellites
        lon0 = (lon0 + lon_eq + 180.0) % 360 - 180.0
        lonc = (lonc + lon_eq + 180.0) % 360 - 180.0

        if np.any((lonc < -180.0) | (lonc > 180.0)):
            raise ValueError

    # load orbit from file
    else:
        if orbit_path is None:
            logging.error('No orbit filename (`orbit_path`) provided.')
            return
        logging.info('Read orbit from "%s"' % orbit_path)
        orbit_data = None
        try:
            times, lon0, lat0, lonc, latc, sza = read_orbit(orbit_path)
        except FileNotFoundError as err:
            logging.error(err)

    # check if any satellite pixels are in COSMO domain
    if not any_pixels_in_domain(lon0, lat0, cfg.DOMAIN, cfg.ROTPOLE):
        logging.warning('All satellite ground pixels are outside COSMO domain.')
        return

    # set overpass time and cosmo file
    if overpass_time is None:
        overpass_time = datetime(2015,1,1) + timedelta(seconds=np.mean(times))

    os.makedirs(output_path, exist_ok=True)
    output_filename = os.path.join(
        output_path,
        '%s_%s_o%04d_l%03d.nc' % (instrument,
            overpass_time.strftime('%Y%m%d%H'),
            orbit_number, lon_eq
        )
    )

    if dont_overwrite and os.path.exists(output_filename):
        logging.warning('Stop, because output file exists')
        return

    # COSMO filename
    if overpass_time.month >= 7:
        cosmo_path = cosmo_path.format(month='07')
    else:
        cosmo_path = cosmo_path.format(month='01')

    cosmo_filename = None if cosmo_path is None else overpass_time.strftime(cosmo_path)

    if cosmo_filename is None or not os.path.exists(cosmo_filename):
        logging.error('COSMO filename "%s" not found. Stop.' % cosmo_filename)
        return

    # MODIS albedo filename
    albedo_filename = None if albedo_path is None else overpass_time.strftime(albedo_path)

    # check for MODIS albedo filename and return if Sentinel-7 CO2
    if instrument == 'Sentinel_7_CO2' and (albedo_filename is None or not os.path.exists(albedo_filename)):
        logging.error('Sentinel_7_CO2 requires an albedo file but "%s" does not exits.' % albedo_filename)
        return

    # COSMO and albedo vars
    if cosmo_vars is None:
        cosmo_vars = cfg.cosmo_vars[instrument]
        cosmo_attrs = cfg.cosmo_attrs

    if albedo_vars is None:
        albedo_vars = cfg.albedo_vars
        albedo_attrs = cfg.albedo_attrs

    # init netCDF file (only if given COSMO file exists)
    logging.info('write netCDF file: %s' % output_filename)
    io.init_ncfile(output_filename, instrument, times, lon0, lat0,
            lonc, latc, sza, orbit_data)

    # process COSMO and MODIS fields
    indices = None
    for name, filename, var_names, attr_names in zip(
            ['COSMO',        'MODIS'],
            [cosmo_filename, albedo_filename],
            [cosmo_vars,     albedo_vars],
            [cosmo_attrs,    albedo_attrs]
    ):
        if filename is None or not os.path.exists(filename):
            continue

        if indices is None:
            # compute pixel indices and weights
            logging.info('compute pixel indices and weights')
            indices = pre_compute_indices(lonc, latc)

        logging.info('process %s fields: ' % name)
        fields, attrs, var_names = read_variables(filename, overpass_time, var_names, attr_names)
        fields = sum_satellite_pixels(fields, lonc, latc, indices=indices)

        for var_name, field, attr in zip(var_names, fields, attrs):
            io.add_variable(output_filename, var_name, field, attr)



    # add uncertainty scenarios to output file
    if instrument == 'Sentinel_7_CO2':
        logging.info('process Sentinel-7 CO2 uncertainty scenarios')

        with Dataset(output_filename) as nc:
            albedo_nir = nc.variables['Albedo_NIR'][:]
            albedo_swir = nc.variables['Albedo_SWIR'][:]

        uXCO2 = epf.epf(sza, albedo_nir, albedo_swir, aod=0.0, cod=0.0,
                cth=10.0, error_type='random', species='CO2')

        # attributes:
        attrs = {
            'standard_name': 'XCO2 random uncertainty using error parametrization formula',
            'long_name':     'XCO2 random uncertainty using error parametrization formula',
            'units':         'ppm',
            'VEG50_uncertainties': '0.5, 0.7, 1.0',
            'scaling_factors':     '0.333, 0.500, 0.667'
        }
        io.add_variable(output_filename, 'uXCO2', uXCO2, attrs)


    if instrument.startswith('Sentinel_7_NO2') or instrument == 'Sentinel_7_CO2':
        logging.info('process Sentinel-7 NO2/CO uncertainty scenarios')

        scenario = cfg.uncertainty_scenarios['Sentinel_7']

        # read total tracer field
        with Dataset(output_filename) as nc:
            NO2 = sum([nc.variables[name][:] for name in scenario['NO2']])
            CO = sum([nc.variables[name][:] for name in scenario['CO']])

        # attributes
        attrs = {
            'uNO2_low': {
                'standard_name': 'NO2 uncertainty (low uncertainty scenario)',
                'long_name':     'NO2 uncertainty (low uncertainty scenario)',
                'units':         'molecules cm-2',
            },
            'uNO2_high': {
                'standard_name': 'NO2 uncertainty (high uncertainty scenario)',
                'long_name':     'NO2 uncertainty (high uncertainty scenario)',
                'units':         'molecules cm-2',
            },
            'uCO_low': {
                'standard_name': 'CO uncertainty (low uncertainty scenario)',
                'long_name':     'CO uncertainty (low uncertainty scenario)',
                'units':         'molecules cm-2',
            },
            'uCO_high': {
                'standard_name': 'CO uncertainty (high uncertainty scenario)',
                'long_name':     'CO uncertainty (high uncertainty scenario)',
                'units':         'molecules cm-2',
            },
        }
        for name in attrs.keys():
            field = NO2 if 'NO2' in name else CO
            u = np.clip(scenario[name][0] * field, scenario[name][1], None)
            io.add_variable(output_filename, name,  u, attrs[name])


    if instrument == 'Sentinel_5':
        logging.info('process Sentinel-5 NO2/CO uncertainty scenarios')

        scenario = cfg.uncertainty_scenarios['Sentinel_5']

        # read total tracer field
        with Dataset(output_filename) as nc:
            NO2 = sum([nc.variables[name][:] for name in scenario['NO2']])
            CO = sum([nc.variables[name][:] for name in scenario['CO']])

        # attributes
        attrs = {
            'uNO2': {
                'standard_name': 'NO2 uncertainty',
                'long_name':     'NO2 uncertainty',
                'units':         'molecules cm-2',
            },
            'uCO': {
                'standard_name': 'CO uncertainty',
                'long_name':     'CO uncertainty',
                'units':         'molecules cm-2',
            },
        }
        for name in attrs.keys():
            field = NO2 if 'NO2' in name else CO
            u = np.clip(scenario[name][0] * field, scenario[name][1], None)
            io.add_variable(output_filename, name,  u, attrs[name])













