

import matplotlib.pyplot as plt
import numpy as np
import pandas



FILENAME_TRUE = '/users/gkuhl/projects/smartcarb/ts/ts_berlin_emissions.dat'

A_PRIORI     = 534.0
A_PRIORI_STD = 262.0

def main(filename_est, filename_true):

    # load emission time series and add a priori uncertainty
    ts = pandas.DataFrame.from_csv(filename_true, sep='\t').resample('H').asfreq()

    ts['a priori'] = pandas.Series(np.full(ts.index.size, A_PRIORI), index=ts.index)
    ts['a priori std'] = pandas.Series(np.full(ts.index.size, A_PRIORI_STD), index=ts.index)

    # average last 3 hours
    ts = ts[(ts.index.hour == 9) | (ts.index.hour == 10) | (ts.index.hour == 11)].resample('D').mean()


    em = pandas.DataFrame.from_csv(filename_est, sep='\t', header=7).resample('D').mean()
    em[em.index > '2015-10-15'].values[:] = np.nan
    print(em.columns)

    em2 = em[em['emissions std'] < A_PRIORI_STD].resample('D').asfreq()

    for month in range(1,13):
        mask = em2.index.month == month
        em2.values[mask,:] = em2[mask].mean().values[np.newaxis]
    em2.values[em2.index.day == 1,:] = np.nan


    # plot data
    ax = ts.plot(y='variable emissions', label='"true" emissions')

    ax = ts.plot(y='a priori', label=' a priori', ax=ax, color='#ff7f0e')
    ax.fill_between(ts.index,
            ts['a priori'] - ts['a priori std'],
            ts['a priori'] + ts['a priori std'],
            facecolor='#ff7f0e', alpha=0.2, interpolate=True,
    )

    ax = em[em['emissions std'] < A_PRIORI_STD].plot(y='emissions', yerr='emissions std',
            ax=ax, marker='.', color='red', ls='', label='a posteriori')
    em2.plot(y='emissions', ax=ax, ls='-', marker='',  color='k', label='a posteriori (mean)')




    ax.set_ylabel('emissions (kg s$^{-1}$)')
    ax.set_xlabel('UTC')
    ax.set_ylim(0,1600)
    plt.legend()
    plt.tight_layout()

    return ts, em, em2

