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

from __future__ import print_function, division
from six import string_types

from datetime import datetime, timedelta
import glob
import os

from xarray import DataArray, Dataset
import numpy as np
import netCDF4
import pandas
import xarray

import c2sm

from smartcarb import config as cfg
from smartcarb import misc


CORRECT_BERLIN = True


def table2filename(ts):
    """
    Use time, satellite, lon_eq and orbit in pandas table
    to create filenames.
    """
    filenames = []
    for time,values in ts.iterrows():
        sat = values.satellite
        lon = values.lon_eq
        orbit = values.orbit
        filename = create_filename(time, starttime=None,
                                   input_files='level2',
                                   instrument=sat, orbit=orbit,
                                   lon_eq=lon)
        filenames.append(filename)

    return np.array(filenames)



def filenames2table(out_filename, filenames):
    columns = ['satellite', 'orbit', 'date', 'lon_eq']
    data = [process_filename(n) for n in filenames]

    d = pandas.DataFrame(data, columns=columns)
    d.to_csv(out_filename, sep='\t', index=False)



def read_level2(filename, co2_noise=None, co2_thr=None, no2_noise=None,
                no2_thr=None, use_constant=False, add_berlin=True,
                seed='orbit', add_co_signal=False, origin='Berlin',):

    data = {}
    dims = ('along', 'across')
    dims2 = ('along', 'across', 'corners')

    satellite, orbit, time, lon_eq =  process_filename(filename)

    lon, lat, lonc, latc = read_ground_pixels(filename, use_xarray=True)
    data['lon'] = lon
    data['lat'] = lat
    data['lonc'] = lonc
    data['latc'] = latc

    data['clouds'] = read_fields(filename, ['CLCT'], use_xarray=True)[0]

    data['psurf'] = read_fields(filename, ['PS'], use_xarray=True)[0]

    # CO2 and NO2
    for tracer, noise, thr in [
        ('CO2', co2_noise, co2_thr),
        ('NO2', no2_noise, no2_thr)
    ]:
        if noise is None or thr is None:
            continue

        attrs = {'noise_level': noise, 'cloud_threshold': thr}

        val = read_total_tracer(filename, tracer, time=time,
                                use_constant_emissions=use_constant,
                                use_xarray=True)

        err = read_tracer_noise(filename, tracer, level=noise,
                                use_epf=True, seed=seed, use_xarray=True)

        is_cloudy = data['clouds'] > thr

        # store noisefree/gapfree data
        data['%s_noisefree'%tracer] = val.copy()

        val[:] = np.where(is_cloudy, np.nan, val)
        err[:] = np.where(is_cloudy, np.nan, err)

        data[tracer] = err + val
        data['%s_noise'%tracer] = err

        data[tracer].attrs.update(attrs)
        data[tracer].attrs.update(val.attrs)
        data[tracer].attrs['standard name'] += ' with random noise'
        data[tracer].attrs['long name'] += ' with random noise'


    # CO2 and NO2 signals of Berlin
    if add_berlin:
        letter = 'B' if origin == 'Berlin' else 'J'
        #if co2_noise is not None:
        signal = 'XCO2_%sC' % letter if use_constant else 'XCO2_%sV' % letter
        data['CO2_true_signal'] = read_fields(filename, [signal], time,
                                              use_xarray=True)[0]

        if no2_noise is not None:
            signal = 'NO2_%sC' % letter if use_constant else 'NO2_%sV' % letter
            data['NO2_true_signal'] = read_fields(filename, [signal], time,
                                                  use_xarray=True)[0]

        if add_co_signal:
            signal = 'CO_%sC' % letter if use_constant else 'CO_%sV' % letter
            data['CO_true_signal'] = read_fields(filename, [signal], time,
                                                 use_xarray=True)[0]

    return Dataset(data, attrs={'satellite': satellite, 'orbit': orbit,
                                'time': time, 'lon_eq': lon_eq})





def get_signal_name(tracer, origin, use_constant_emissions=False):
    """
    """
    if use_constant_emissions:
        suffix = 'BC' if origin == 'Berlin' else 'JC'
    else:
        suffix = 'BV' if origin == 'Berlin' else 'JV'

    if tracer == 'CO2':
        tracer = 'XCO2'

    return '%s_%s' % (tracer, suffix)



def read_emissions(origin, time, stop=None, use_constant=False, overpass=11):

    column = '%s (Mt/y)' % origin

    filename = os.path.join(cfg.DATA_PATH, '..', 'flux_inversion',
                            'ts_CO2_emissions.dat')
    d = pandas.read_csv(filename, sep='\t', index_col=0, parse_dates=True)

    if use_constant:
        d[:] = {
            'Boxberg':       19.055633, # Mt/yr
            'Janschwalde':   33.322991,
            'Lippendorf':    15.256415,
            'Patnow':         5.826520,
            'Schwarze Pumpe': 8.211211,
            'Turow':          8.716601,
            'Berlin':        16.85#23.805240
        }[origin]

    if time == 'raw':
        return d[column]

    if time == 'all':
        return d[column], np.nan
    elif time == 'overpass':
        return d[d.index.hour == overpass][column], np.nan
    else:
        if stop is None:
            v = d[column][d.index == time]
        else:
            v = d[column][(time <= d.index) & (d.index <= stop)]

        return np.mean(v), np.std(v)



def read_wind_profile(origin, time, path=cfg.WIND_PATH):

    fn1 = os.path.join(path, 'ts_U_%s.dat' % origin)
    fn2 = os.path.join(path, 'ts_V_%s.dat' % origin)

    with open(fn1) as f1, open(fn2) as f2:

        h = np.array(next(f1).split('\t'), 'f4')
        h = np.array(next(f2).split('\t'), 'f4')

        for line in f1:
            line = line.split('\t')
            if time == datetime.strptime(line[0], '%Y-%m-%d %H:%M:%S'):
                u = np.array(line[1:], 'f4')
                break

        for line in f2:
            line = line.split('\t')
            if time == datetime.strptime(line[0], '%Y-%m-%d %H:%M:%S'):
                v = np.array(line[1:], 'f4')
                break

    return h, u, v


def process_filename(filename):
    #Sentinel_7_CO2_2015051511_o2009_l270.nc
    filename = os.path.basename(filename)
    filename = os.path.splitext(filename)[0]
    filename = filename.split('_')
    if filename[1] == '5':
        sat, n, date, orbit, lon = filename[:5]
        sat = '_'.join([sat, n])
    else:
        sat, n, t, date, orbit, lon = filename[:6]
        sat = '_'.join([sat, n, t])

    orbit = int(orbit[1:])
    date = datetime.strptime(date, '%Y%m%d%H')
    lon = int(lon[1:])

    return sat, orbit, date, lon


def filename2orbit(filename):
    sat, orbit, date, lon = process_filename(filename)
    return orbit


def filename2label(filename):
    sat, orbit, date, lon = process_filename(filename)

    sat = 'CO$_2$ satellite'
    label = '%s ($\lambda_{eq}$: %d°, orbit: %d, %s)' % (sat, lon/100, orbit,
            date.strftime('%Y-%m-%d %H:%M UTC'))

    return label


def filename2time(filename):
    #Sentinel_7_CO2_2015051511_o2009_l270.nc
    sat, orbit, date, lon = process_filename(filename)
    return date


def s7_to_s5(filename, hours=[9,8]):
    """
    Try to find Sentinel-5 Level-2 data for given Sentinel-7 file.
    """
    path = os.path.dirname(filename).replace('Sentinel_7_CO2', 'Sentinel_5')
    sat, orbit, date, lon = process_filename(filename)

    pattern = os.path.join(path,
        'Sentinel_5_%s??_o*_l0000.nc' % (date.strftime('%Y%m%d'))
    )
    #print(glob.glob(pattern))

    for hour in hours:
        pattern = os.path.join(path,
            'Sentinel_5_%s%02d_o*_l0000.nc' % (date.strftime('%Y%m%d'), hour)
        )
        fnames = sorted(glob.glob(pattern), reverse=True)

        if fnames:
            return fnames[0]
    else:
        raise ValueError('No Sentinel-5 filename found')






def str2char(s):
    return np.array(list(s), 'c')


def init_netcdf(nc, rlon, rlat, lon=None, lat=None, pollon=cfg.pollon, pollat=cfg.pollat,
        attrs=None):
    """\
    Init netCDF4 file with cosmo rotated pole grid
    """
    if attrs is not None:
        nc.setncatts(attrs)

    # dimensions
    nc.createDimension('time')
    nc.createDimension('rlat', rlat.size)
    nc.createDimension('rlon', rlon.size)

    # time variable
    nc_time = nc.createVariable('time', 'f8', ('time',))
    nc_time.setncattr('units', str2char('seconds since 2015-01-01 00:00'))
    nc_time.setncattr('calendar', str2char('proleptic_gregorian'))

    # rlon
    nc_rlon = nc.createVariable('rlon', 'f8', ('rlon',))
    nc_rlon[:] = rlon
    nc_rlon.setncattr('axis', str2char('X'))
    nc_rlon.setncattr('units', str2char('degrees'))
    nc_rlon.setncattr('standard_name', str2char('longitude'))

    # rlat
    nc_rlat = nc.createVariable('rlat', 'f8', ('rlat',))
    nc_rlat[:] = rlat
    nc_rlat.setncattr('axis', str2char('Y'))
    nc_rlat.setncattr('units', str2char('degrees'))
    nc_rlat.setncattr('standard_name', str2char('latitude'))

    # rotated pole
    nc_pole = nc.createVariable('rotated_pole', 'c', ('time','rlat','rlon'))
    nc_pole.setncattr('grid_mapping_name', str2char('rotated_latitude_longitude'))
    nc_pole.setncattr('grid_north_pole_latitude', pollat)
    nc_pole.setncattr('grid_north_pole_longitude', pollon)
    nc_pole.setncattr('north_pole_grid_longitude', 0.0)

    # longitude
    if lon is not None:
        nc_lon = nc.createVariable('lon', 'f8', ('rlat', 'rlon'))
        nc_lon[:] = lon
        nc_lon.setncattr('axis', str2char('X'))
        nc_lon.setncattr('units', str2char('degrees'))
        nc_lon.setncattr('standard_name', str2char('longitude'))

    # latitude
    if lat is not None:
        nc_lat = nc.createVariable('lat', 'f8', ('rlat', 'rlon'))
        nc_lat[:] = lat
        nc_lat.setncattr('axis', str2char('Y'))
        nc_lat.setncattr('units', str2char('degrees'))
        nc_lat.setncattr('standard_name', str2char('latitude'))



def append_variable(nc, name, values, attrs=None):
    """\
    Append variable (name, values) to netCDF file.
    """
    if attrs is None:
        attrs = {}

    var = nc.createVariable(name, values.dtype, ('time', 'rlat', 'rlon'), zlib=True)
    var[0] = values

    # add or update attributes
    for key, value in attrs.items():
        if isinstance(value, string_types):
            var.setncattr(key, str2char(value))
        else:
            var.setncattr(key, value)


def get_slice(origin, lon=None, lat=None, area=100):
    """\
    e.g. Berlin, area=40
    e.g. Janschwalde, area=30
    """

    if lon is None or lat is None:
        lon, lat = c2sm.transform.rotpole2wgs(cfg.rlon2, cfg.rlat2, cfg.pollon, cfg.pollat)

    if origin in cfg.CITIES:
        lon_o, lat_o = cfg.CITIES[origin]
    else:
        lon_o, lat_o = cfg.POINT_SOURCES[origin]

    i,j = misc.nearest_index(lon_o, lat_o, lon, lat)

    return slice(i-area//2,i+area//2+1), slice(j-area//2,j+area//2+1)


def list_level2_filenames(filename, path='.'):

    fnames = []

    with open(filename) as list_file:
        for line in list_file:
            if line.startswith('#'):
                continue

            date, sat = line.strip().split()
            lon = {'1': 0, '2': 90, '3': 180, '4': 270}[sat]
            date = date.replace('-', '')

            pattern = os.path.join(path, '*_%s??_o????_l%03d.nc' % (date, lon))

            fnames += glob.glob(pattern)

    return fnames





def create_filename(time, starttime=None, input_files='3D', instrument='', orbit=0,
        lon_eq=0, path=cfg.DATA_PATH, use_restart_structure=False,
    ):
    """\
    Create filename for netCDF file.
    """
    if starttime is None:
        starttime = time2starttime(time)

    if input_files in ['3Dc', '3D']:
        hstart = (time - starttime).days // 3 * 72
        hstop = hstart + 72

        if use_restart_structure:
            path = os.path.join(
                    path,
                    '%s_%d_%d' % (starttime.strftime('%Y%m%d%H'), hstart, hstop),
                    'cosmo_output'
            )
        else:
            path = os.path.join(path, starttime.strftime('cosmo_%Y-%m'))

        if input_files == '3D':
            filename = os.path.join(
                    path, time.strftime('lffd%Y%m%d%H.nc')
            )
        else:
            filename = os.path.join(
                    path, starttime.strftime('lffd%Y%m%d%Hc.nc')
            )

    elif input_files == '2D':
        path = os.path.join(path, starttime.strftime('cosmo2D_%Y-%m'))
        filename = os.path.join(
                path, time.strftime('cosmo_2d_%Y%m%d%H.nc')
        )

    elif input_files == 'level2':
        # Level-2 data for instrument, orbit and lon_eq

        if not isinstance(orbit, string_types):
            orbit = '%04d' % orbit
        if not isinstance(lon_eq, string_types):
            lon_eq = '%04d' % lon_eq

        pattern = '{0:s}_%Y%m%d??_o{1:s}_l{2:s}.nc'.format(instrument, orbit, lon_eq)
        pattern = os.path.join(path, 'level2', instrument, time.strftime(pattern))

        filename = glob.glob(pattern)

        if not filename:
            print(pattern)
        if len(filename) > 1:
            print('Use first file of %d files.' % len(filename))
        filename = filename[0]

    return filename



def read_3D_for_level2(level2_filename, fields, type_='3D', path=None,
        correct_berlin=CORRECT_BERLIN, slice_=Ellipsis):
    """
    Read 3D data for Level-2 data.
    """
    time = process_filename(level2_filename)[2]
    starttime = time2starttime(time)

    if path is None: # TODO: move to config
        path = starttime.strftime('/project/d73/output/cosmo_%Y-%m/')

    filename = create_filename(time, starttime, type_, path=path)

    return read_fields(filename, fields, time, correct_berlin, slice_)


def read_ground_pixels(filename, slice_=Ellipsis, use_xarray=False):
    """Read center and corners of ground pixels."""
    return read_fields(filename, [
            'longitude', 'latitude', 'longitude_corners', 'latitude_corners'
        ], slice_=slice_, use_xarray=use_xarray)



def read_total_tracer(filename, tracer, time=None,
                      correct_berlin=CORRECT_BERLIN, clip_data=True,
                      slice_=Ellipsis, use_constant_emissions=False,
                      use_xarray=False):
    """\
    Read total tracer field (CO2, NO2 or CO).
    """
    if time is None:
        time = filename2time(filename)

    suffix = '_BC' if use_constant_emissions else '_BV'

    if tracer in ['CO2', 'XCO2']:
        fields = ['XCO2%s' % suffix, 'XCO2_A', 'XCO2_JV', 'XCO2_RA', 'XCO2_GPP', 'XCO2_BG']
        bv,a,jv,ra,gpp,bg = read_fields(filename, fields, time, correct_berlin,
                                        slice_, use_xarray=use_xarray)

        field = bv + a + jv + ra - gpp + bg


        if use_xarray:
            attrs = {
                'standard name': 'CO2_column-averaged_dry-air_mole_fraction',
                'long name': 'CO2_column-averaged_dry-air_mole_fraction',
                'units': 'ppm'
            }
            field.attrs.update(attrs)

        if clip_data:
            field[:] = np.where((field < 0) | (field > 1e5), np.nan, field)

        return field

    else:
        if tracer == 'NO2':
            fields = ['NO2%s' % suffix, 'NO2_A', 'NO2_JV', 'NO2_BG']
            bv,a,jv,bg = read_fields(filename, fields, time, correct_berlin,
                                     slice_, use_xarray=use_xarray)
            field =  bv + a + jv + bg

        if tracer == 'CO':
            fields = ['CO%s' % suffix, 'CO_A', 'CO_JV', 'CO_BG']
            bv,a,jv,bg = read_fields(filename, fields, time, correct_berlin,
                                     slice_, use_xarray=use_xarray)
            field = bv + a + jv + bg

        if use_xarray:
            attrs = {
                'standard name': '%s_vertical_column_density' % tracer,
                'long name': '%s_vertical_column_density' % tracer,
                'units': 'molecules cm-2'
            }
            field.attrs.update(attrs)

        return field



def read_tracer_noise(filename, tracer, use_epf=True, level=0.7, slice_=Ellipsis,
                      seed='orbit', use_xarray=False):
    """
    tracer in {'CO2', 'NO2', 'CO', 'uXCO2', 'uNO2_low', 'uNO2_high',
               'uCO_low', 'uCO_high'}
    level in {0.5, 0.7, 1.0, 'low', 'high'}
    """
    sat,orbit,_,lon_eq = process_filename(filename)

    if sat == 'Sentinel_5':
        utracer = 'uNO2'
    else:
        if tracer == 'CO2':
            utracer = 'uXCO2'
        elif tracer == 'NO2':
            utracer = 'uNO2_%s' % level
        else:
            utracer = 'uCO_%s' % level

    # read uncertainty (TODO/FIXME: correct uNO2 for Jan/Jul)
    u = read_fields(filename, [utracer], slice_=slice_, use_xarray=use_xarray)[0]

    if not use_xarray:
        u[u.mask] = np.nan
        u = u.data

    # seed based on orbit and lon_eq for same noise patterns with each species
    if seed == 'orbit':
        np.random.seed(10000 * lon_eq + orbit)
    else:
        np.random.seed(seed)

    noise = np.random.randn(*u.shape).astype(u.dtype)

    if tracer == 'CO2':
        if use_epf:
            noise = level * noise * u / 1.50
        else:
            noise = level * noise

    elif tracer == 'NO2' and use_epf:
        cc = read_fields(filename, ['CLCT'], slice_=slice_)[0]

        # use Wenig et al. (2008) with `uNO2` instead of 1.5e15
        noise = (1.0 + 3.0 * cc) * u * noise

    else:
        noise = u * noise

    if use_xarray:
        if tracer == 'CO2':
            name = '%s random noise (using %s and %.1f ppm)' % (tracer, utracer, level)
        else:
            name = '%s random noise (using %s)' % (tracer, utracer)

        noise.attrs['standard name'] = name
        noise.attrs['long name'] = name
        noise.attrs['units'] = 'ppm' if tracer == 'CO2' else 'molecules cm-2'

    return noise





def read_fields(filename, tracers, time=None, correct_berlin=CORRECT_BERLIN,
        slice_=Ellipsis, use_xarray=False):
    """\
    Read fields from netCDF file. If correct_berlin is True,
    correct '*_BV' tracers for January and July
    by scaling '*_B0' with 0.55 which corrects the too high
    emissions in January and July in the SMARTCARB simulations.
    """
    with netCDF4.Dataset(filename) as nc:
        fields = []

        for tracer in tracers:
            if correct_berlin and tracer.endswith('_BV'):

                field = nc.variables[tracer]
                attrs = dict((k, field.getncattr(k)) for k in field.ncattrs())

                label = tracer.split('_')[0]

                berlin_tracers = ['%s_B%d' % (label, i) for i in range(3)]
                b0,b1,b2 = read_fields(filename, berlin_tracers, time,
                        correct_berlin, slice_, use_xarray=use_xarray)

                if time.month in [1,7]:
                    field = 0.55 * b0 + b1 + b2
                else:
                    field = b0 + b1 + b2

                if use_xarray:
                    field.attrs.update(attrs)

            else:
                field = nc.variables[tracer]

                if use_xarray:
                    attrs = dict((k, field.getncattr(k)) for k in field.ncattrs())
                    field = DataArray(field[:], dims=field.dimensions, attrs=attrs)

                if slice_ is Ellipsis or isinstance(slice_, slice):
                    field = field[slice_]
                else:
                    # use slice starting at last dimension (instead of first)
                    ndim = np.ndim(field)

                    if len(slice_) == ndim:
                        field = field[slice_]
                    else:
                        s = [Ellipsis] * (ndim - len(slice_)) + list(slice_)
                        field = field[s]


            # correct _BC fields
            if correct_berlin and tracer.endswith('_BC') and time.month in [1,7]:
                field *= 0.708

            field = np.squeeze(field)
            fields.append(field)

    return fields




def read_field(time, starttime, keys, indices=Ellipsis, input_files='3D',
        path='.', instrument='', orbit=0, lon_eq=0
    ):
    if indices is Ellipsis:
        indices = [Ellipsis for _ in keys]


    if input_files in ['3Dc', '3D']:
        hstart = (time - starttime).days // 3 * 72
        hstop = hstart + 72

        folder = '%s_%d_%d' % (starttime.strftime('%Y%m%d%H'), hstart, hstop)

        if input_files == '3D':
            filename = os.path.join(
                    path, folder, 'cosmo_output',
                    time.strftime('lffd%Y%m%d%H.nc')
            )
        else:
            filename = os.path.join(
                    path, folder, 'cosmo_output',
                    starttime.strftime('lffd%Y%m%d%Hc.nc')
            )

    elif input_files == '2D':
        path = os.path.join(path, starttime.strftime('cosmo2D_%Y-%m'))
        filename = os.path.join(
                path, time.strftime('cosmo_2d_%Y%m%d%H.nc')
        )

    # read Level-2 data for instrument, orbig and lon_eq
    elif input_files == 'level2':

        filename = '{0:s}_%Y%m%d??_o{1:04d}_l{2:03d}.nc'.format(instrument, orbit, lon_eq)
        filename = os.path.join(path, 'level2', instrument, time.strftime(filename))

        filename = glob.glob(filename)
        if len(filename) > 1:
            print('Use first file of %d files.' % len(filename))
        filename = filename[0]


    print(filename)
    with netCDF4.Dataset(filename) as nc:
        data = [nc.variables[key][index] for key, index in zip(keys, indices)]

    return data[0] if len(data) == 1 else data


def time2starttime(time):
    if time >= datetime(2015,7,1):
        starttime = datetime(2015,7,1)
    else:
        starttime = datetime(2015,1,1)
    return starttime



def read_signal_and_background(filename, tracer, time=None,
        correct_berlin=CORRECT_BERLIN, slice_=Ellipsis, add_biosphere=False):

    signal = read_fields(filename, [tracer], time=time,
            correct_berlin=correct_berlin, slice_=slice_)[0]
    totals = read_total_tracer(filename, tracer.split('_')[0], time=time,
            correct_berlin=correct_berlin, slice_=slice_)

    bg = totals - signal

    if add_biosphere:
        gpp, ra = read_fields(filename, ['XCO2_GPP', 'XCO2_RA'], time=time,
                              slice_=slice_)
        bio = ra - gpp

        return signal, bg, bio

    return signal, bg





if __name__ == '__main__':
    pass


