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

from datetime import datetime, timedelta
import os

from matplotlib.dates import DateFormatter
import matplotlib.pyplot as plt
import numpy as np
import pandas
import scipy.stats
import scipy.ndimage as ndi
import scipy.optimize as opt
import statsmodels.api as sm

import doas

from smartcarb import io
from smartcarb import misc

import dplume

# kg/s -> Mt/y (Tg/y)
kgs_to_Mty = 1e3 * (365 * 24 * 60 * 60) / 1e12



def compute_emissions(nc):

    cats = [
       'CO2_point_snap1_major', # this was counted twice in Jan and Jul
       'CO2_point_snap1_minor',
       'CO2_point_snap2_minor',
       'CO2_point_snap3_minor',
       'CO2_point_snap4_minor',
       'CO2_point_snap5_minor',
       'CO2_point_snap34_minor',
       'CO2_point_snap8_minor',
       'CO2_point_snap9_minor',
       'CO2_area_construction',
       'CO2_line_aircraft',
       'CO2_point_aircraft',
       'CO2_area_aircraft',
       'CO2_area_forestry',
       'CO2_area_offroad',
       'CO2_line_shipping',
       'CO2_point_shipping',
       'CO2_line_train',
       'CO2_area_heating',
       'CO2_line_majorroads',
       'CO2_area_minorroads'
    ]

    total = 0.0
    for cat in cats:
        total += np.sum(nc.variables[cat][:] * 1235941.1)

    return total



def get_constant_emissions(time, use_constant=True):
    if not use_constant and time.year == 2015 and time.month in [1,7]:
        return 754.86
    else:
        return 534.30



def simple_leastsq(filename, error_std=0.7, seed='orbit',
        background_mode='model', use_constant_emissions=False
    ):
    """\
    filename is either a file with 2d COSMO field or Level-2 data

    y[i,j] = x * H[i,j] + b

    - y: XCO2 measurements in plume
    - H: XCO2 from Berlin tracer (e.g. XCO2_BC) normalized by emissions
    - b: background (e.g. mean of "background" tracers in COSMO)

    The plume is defined where 'plume_tracer' > 'plume_thr'.

    If `use_const_emissions` use XCO2_BV in observation vector y.

    Names for H:
    """
    # tracer fields (synthetic observations)
    time = io.filename2time(filename)

    data = io.read_level2(filename, co2_noise=error_std, co2_thr=0.01,
                          use_constant=use_constant_emissions,
                          add_berlin=True)

    # signals (constant and variable)
    bv, bc = io.read_fields(filename, ['XCO2_BV', 'XCO2_BC'], time, correct_berlin=True)


    # define plume based on signal and cloud threshold
    plume = dplume.create_true_plume(data, 'CO2', 'Berlin',
                                     signal_thr=0.05, cloud_thr=1.0)
    plume = np.array(plume) & np.array(data.clouds <= 0.01)

    # remove data gaps (high uncertainty)
    plume &= np.array(np.isfinite(data.CO2) & np.isfinite(data.clouds))


    if np.sum(plume) <= 1:
        return np.nan, np.nan, np.sum(plume)


    # background
    b = np.array(data.CO2 - data.CO2_noise - data.CO2_true_signal)[plume]

    if background_mode == 'model':
        bg_std = 0.0
    elif background_mode == 'mean':
        bg_std = b.std()
        b = b.mean()
    else:
        raise ValueError('unknown mode for background estimation')

    # observation vector
    y = np.array(data.CO2)[plume]

    # footprint
    #if use_constant_emissions:
    #    H = bc[plume] / 534.30
    #else:
    #    H = bv[plume] / 534.30
    H = bc[plume] / 534.30

    # Se, xa and Sa
    if error_std == 0.0:
        Se = np.nan
    else:
        Se = np.full(y.shape, error_std**2)

    # forward model
    f = lambda x,_: H * x / kgs_to_Mty + b

    x0 = np.array([534.0]) * kgs_to_Mty
    x, res = doas.solver.gauss_newton(x0, y, f, Se=Se)

    return x[0], np.sqrt(float(res['S'])), np.sum(plume)



def main(out_filename, error_std=0.7, nmin=0,
         use_constant_emissions=False, background_mode='model'):

    filename = "/project/kug/smartcarb/results/Berlin/BerlinPlume.dat"
    data = pandas.read_csv(filename, sep='\t', index_col=2, parse_dates=True)

    level2_filenames = [
        io.create_filename(data.index[i], None, 'level2', 'Sentinel_7_CO2',
                           data.orbit[i], data.lon_eq[i])
        for i in range(data.shape[0])
        if data['pixels (<1%)'][i] > nmin
    ]

    em = pandas.read_csv('/project/smartcarb/flux_inversion/ts_berlin_emissions.dat',
                         delimiter='\t', index_col=0, parse_dates=True)
    em = em['CO2_BV_fix']

    with open(out_filename, 'w') as out_file:
        out_file.write('\t'.join([
            'UTC',
            'lon_eq',
            'orbit',
            'number of pixels',
            'CO2 emissions (Mt/yr)',
            'CO2 uncertainty (Mt/yr)',
            'CO2 true emissions (Mt/yr)',
        ]))
        out_file.write('\n')

        for fn in level2_filenames:

            _, orbit, time, lon_eq = io.process_filename(fn)

            x, x_std, N = simple_leastsq(
                fn, error_std=error_std,
                use_constant_emissions=use_constant_emissions,
                background_mode=background_mode
            )

            if use_constant_emissions:
                true_em = 16.849540
            else:
                true_em = em[time]

            print(time, N, x, true_em)

            out_file.write('\t'.join([
                str(time),
                str(lon_eq),
                str(orbit),
                str(N),
                '%10.3f' % x,
                '%10.3f' % x_std,
                '%10.3f' % true_em
            ]))
            out_file.write('\n')





def read_results(filename, nsat=1):
    ts = pandas.read_csv(filename, delimiter='\t', index_col=0,
                         parse_dates=True, na_values='%10.3f' % np.nan)
    ts = misc.filter_dataset(ts, nsat)

    return ts


def plot_scatter(filename, nsat=1):
    ts = read_results(filename, nsat)

    est = ts['CO2 emissions (Mt/yr)']
    true = ts['CO2 true emissions (Mt/yr)']

    fig = plt.figure(figsize=(5,5))
    ax = fig.add_subplot(111, aspect='equal')

    ax.plot(true, est, 'o')

    ax.set_xlim(xmin=0)
    ax.set_ylim(ymin=0)


def sqrt_regression(n,y):

    x = 1.0 / np.sqrt(n)

    def model(c, b=None):
        return np.polyval(c,x)

    c0 = np.array([1.0, 0.0])
    c, res = doas.solver.gauss_newton(c0, y, model)
    c_std = np.sqrt(res['S'].diagonal())

    return c, c_std



def plot_vs_number(filename, nsat):

    fig = plt.figure()
    ax = fig.add_subplot(111)

    ts = read_results(filename, nsat)

    est = ts['CO2 emissions (Mt/yr)']
    true = ts['CO2 true emissions (Mt/yr)']

    n = ts['number of pixels']

    valids = n > 5
    n = n[valids]
    y = ts['CO2 uncertainty (Mt/yr)'][valids]

    x = 1.0 / np.sqrt(n)
    c, c_std = sqrt_regression(n, y)


    n_fit = np.arange(0, n.max())
    x_fit = 1.0 / np.sqrt(n_fit)

    y_fit = np.polyval(c, x_fit)
    u_fit = np.sqrt(c_std[0]**2 * x_fit**2 + c_std[1]**2)

    ax.fill_between(n_fit, y_fit - 3.0 * u_fit, y_fit + 3.0 * u_fit)
    ax.plot(n_fit, y_fit, 'k-', label='fit ($\pm$3$\sigma$)')
    ax.plot(n, y, 'o', ms=3, label='data', color='#ff7f0e')

    # compute number of pixels required for a given precision
    line = []
    for thr in [10,25,50,100]: # precision (in %)
    #for thr in [5,10]: # precision (in %)
        thr = thr / 100 * 16.85
        v = (thr - c[1]) / c[0]

        v_std = (-(thr - c[1]) / c[0]**2)**2 * c_std[0]**2
        v_std += (-1.0 / c[0])**2 * c_std[1]**2
        v_std = np.sqrt(v_std)

        vmin = 1.0 / (v - v_std)**2
        vmax = 1.0 / (v + v_std)**2
        vmed = 1.0 / v**2

        print('%d(+%d,-%d)' % (round(vmed), round(abs(vmed-vmin)),
                               round(abs(vmed-vmax))))

        line.append('%d-%d (%d)' % (round(vmin), round(vmax), round(vmed)))

        #ax.plot([vmin,vmax], [thr,thr], 'ro-')

    #print('\t'.join(line))


    ax.set_xlabel('number of pixels')
    ax.set_ylabel('CO$_2$ uncertainty (Mt/yr)')

    ax.set_xlim(0,1500)
    ax.set_ylim(0,50)

    ax.grid(True)

    plt.legend()
    plt.tight_layout()



def calculate_scatter(filename, nsat=1, nmin=0, nmax=np.inf):
    ts = read_results(filename, nsat)

    n = ts['number of pixels']
    valids = (nmin <= n) & (n < nmax)

    est = ts['CO2 emissions (Mt/yr)'][valids]
    unc = ts['CO2 uncertainty (Mt/yr)'][valids]

    true = ts['CO2 true emissions (Mt/yr)'][valids]

    if sum(valids) >= 10:
        scatter = np.std(est - true)
    else:
        scatter = np.nan

    return np.median(unc), scatter, 100 * scatter / true.mean(), sum(valids)


def boxplot(filenames, nsat=6, nmax=100, sample='all', indx=None,
            origin='Berlin', use_constant=False, vmin=-20, vmax=20):

    em10, em11, em_mean = read_true_emissions(use_constant, origin)

    data = []
    unc = []
    for filename in filenames:
        ts = read_results(filename, nsat)
        true = em_mean[ts.index]

        try:
            good = np.isfinite(ts['CO2 emissions (Mt/yr)'])
            good &= ts['CO2 uncertainty (Mt/yr)'] < 0.5*16.85
            #good &= ts['number of pixels'] > 100

            d = (ts['CO2 emissions (Mt/yr)'] - true)[good]
            u = ts['CO2 uncertainty (Mt/yr)'][good]
        except KeyError:
            good = np.isfinite(ts['est. CO2 emissions (Mt/yr)'])
            if 'failed' in ts:
                good &= ts.failed == 0

            d = (ts['est. CO2 emissions (Mt/yr)'] - true)[good]
            u = ts['est. CO2 emissions uncertainty (Mt/yr)']

        data.append(d)
        unc.append(u)


    if sample == 'same':
        ind = sorted((data[0].index) & set(data[1].index) & set(data[2].index))
        data = [d[ind] for d in data]

    if indx is not None:
        data = [d[indx] for d in data]

    line = '\t'.join(['#plumes', 'MB', 'MB', 'SD', 'SD',
                      'RMSD', 'RMSD', 'unc', 'unc'])
    print(line)

    for d,u in zip(data, unc):
        if origin == 'Berlin':
            t = 100 / em_mean.mean()

        n = len(d)
        mb  = np.mean(d)
        std = np.std(d)
        rms = np.sqrt(mb**2 + std**2 / n)

        print('%5d' % n, end='\t')
        print('%5.1f\t%5.1f' % (mb, mb*t), end='\t')
        print('%5.1f\t%5.1f' % (std, std*t), end='\t')
        print('%5.2f\t%5.1f' % (rms, rms*t), end='\t')
        print('%5.1f\t%5.1f' % (np.mean(u), np.mean(u)*t), end='\t')
        print()

        #for n in range(1,7):
        #    print('%d: %5.2f' % (n, total * np.sqrt(6/n)))
        #print()

    fig = plt.figure(figsize=(4,4))
    ax = fig.add_subplot(111)

    for i,n in enumerate(data):
        ax.text(i+1,vmax+0.1, str(len(n)), ha='center', va='bottom')


    ax.boxplot(data, showfliers=True, whis=[5,95], bootstrap=5000,
               labels=['0.5', '0.7', '1.0'], meanline=True, showmeans=True)

    ax.set_xlabel('$\sigma_{VEG50}$ (ppm)')
    ax.set_ylabel('estimated minus true emissions (Mt/yr)')

    ax.set_ylim(vmin,vmax)
    ax.grid(True)

    plt.tight_layout()

    return data


def read_true_emissions(use_constant=False, origin='Berlin'):
    em10 = io.read_emissions(origin, 'overpass', use_constant=use_constant,
                            overpass=10)[0]
    em11 = io.read_emissions(origin, 'overpass', use_constant=use_constant,
                            overpass=11)[0]

    em_mean = pandas.Series(0.5 * (em10.values + em11.values), index=em11.index)

    return em10, em11, em_mean


def plot_ts(filename, nsat=1, use_constant=False, filter_data=True, label=None,
           origin='Berlin', nmax=-1):

    ts = read_results(filename, nsat)
    lon_eqs = misc.lon_eq_by_nsat(nsat)
    em10, em11, em_mean = read_true_emissions(use_constant, origin)
    true = em_mean[ts.index]

    try: # model
        #good = ts['CO2 emissions (Mt/yr)'] >= 0.0
        good = np.isfinite(ts['CO2 emissions (Mt/yr)'])
        #good &= ts['number of pixels'] > 99
        good &= ts['CO2 uncertainty (Mt/yr)'] < 0.5*16.85

        est = ts['CO2 emissions (Mt/yr)'][good]
        est_std = ts['CO2 uncertainty (Mt/yr)'][good]
        true = true[good]

    except KeyError: # data-driven
        #good = ts['est. CO2 emissions (Mt/yr)'] >= 0.0
        good = np.isfinite(ts['est. CO2 emissions (Mt/yr)'])

        if 'plume size' in ts:
            good &= ts['plume size'] > nmax

        if 'failed' in ts:
            good &= ts.failed == 0
        #good &= np.abs(ts['est. CO2 emissions (Mt/yr)'] - true) <= 20

        #good &= ts['plume height (m)'] < 500.
        #good &= ts['est. wind (m/s)'] < 20.
        #good &= ts['est. wind (m/s)'] > 2.

        est = ts['est. CO2 emissions (Mt/yr)'][good]
        est_std = ts['est. CO2 emissions uncertainty (Mt/yr)'][good]
        true = true[good]


    if origin == 'Berlin':
        print(true.mean(), em_mean.mean())
        t = 100 / em_mean.mean()

    diff = est - true
    MB = np.mean(diff)
    SD = np.std(diff)
    RMSD = np.sqrt(MB**2 + SD**2 / good.sum())

    print('N:    %d' % good.sum())
    print('MB:   %5.2f Mt, %.1f%%' % (MB, t*MB))
    print('SD:   %5.2f Mt, %.1f%%' % (SD, t*SD))
    print('RMSD: %5.2f Mt, %.1f%%' % (RMSD, t*RMSD))

    print()
    print('STD (mean):', np.mean(est_std))
    print('STD (median):', np.median(est_std))
    print()


    fig = plt.figure(figsize=(8.4, 4.8))
    ax = fig.add_subplot(111)

    ax.plot(em_mean.index, em_mean, 'k-',
            label='true emissions (10-11 UTC)')
    ax.fill_between(em10.index, em10, em11, color='silver')

    times = ts.index[good]
    try:
        values = ts['CO2 emissions (Mt/yr)'][good]
        errors = ts['CO2 uncertainty (Mt/yr)'][good]
    except:
        values = ts['est. CO2 emissions (Mt/yr)'][good]
        errors = ts['est. CO2 emissions uncertainty (Mt/yr)'][good]

    for i, lon_eq in enumerate(lon_eqs, 1):
        sat = ts[good].lon_eq == lon_eq
        print(i, lon_eq, sat.sum())
        ax.errorbar(times[sat], values[sat], errors[sat], ls='', marker='o',
                    label='#%d (%d plumes)' % (i, sum(sat)))

    if label is not None:
        for i,(t,v) in enumerate(zip(times,values)):
            s = '%.2g' % ts[good][label][i]
            ax.text(t,v,s)

    ax.xaxis.set_major_formatter(DateFormatter('%b'))

    ax.set_xlim('2015-01-01', '2015-12-31')
    if origin == 'Berlin':
        ax.set_ylim(-10,60)
    else:
        ax.set_ylim(ymin=0)

    ax.set_xlabel('2015')
    ax.set_ylabel('CO$_2$ emissions (Mt/yr)')
    ax.grid(True)


    plt.legend(loc=1, ncol=3)
    plt.tight_layout()

    return ts[good]



