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

import os
import sys

from datetime import datetime, timedelta
from netCDF4 import Dataset
import calendar

import matplotlib.pyplot as plt

import numpy as np
import scipy.stats as stats
import scipy.ndimage as ndi
import pandas

import smartcarb as sc
import config as cfg



def n_valids(v, vmin, vmax):
    is_valid = (vmin <= v) & (v < vmax)
    return is_valid.sum()



def plot(filename, tracer):

    if tracer == 'CO2':
        noises = [0.5, 0.7, 1.0]
        ymax = 200 if 'berlin' in filename else 1200
    elif tracer == 'NO2':
        noises = [1e15, 2e15]
        ymax = 700 if 'berlin' in filename else 2200
    elif tracer == 'CO':
        noises = [4e17]
        ymax = 200


    ts = pandas.DataFrame.from_csv(filename, sep='\t').resample('D').mean()
    #ts.values[ts.index > '2015-10-15'] = np.nan

    data = dict()
    valid_days = dict()

    for month in np.arange(1,13):
        mdata = []

        mask = (ts.index.month == month)
        pixels = np.array([ts[mask]['cloud-free pixels %g' % noise].values
            for noise in noises])

        counter = 0

        for day in np.arange(1, calendar.monthrange(2015,month)[1]+1):

            try:
                v = pixels[:,day-1]
            except IndexError:
                continue

            if np.all(np.isfinite(v)):
                counter += 1
                if np.any(v > 0):
                    mdata.append([day, v])

        data[month] = mdata
        valid_days[month] = counter


    most_days = max(len(data[k]) for k in data)

    if most_days == 0:
        print('no days')

    # make figure
    plt.figure(figsize=[12.8, 4.8])

    for i,month in enumerate(np.arange(1,13)):

        days = np.array([v[0] for v in data[month]])
        x = np.arange(0,most_days)[:len(days)] + i * most_days

        # add some text
        xx = most_days * month - np.floor(most_days / 2)
        yy = 0.0
        plt.text(xx,yy, datetime(2015,month,1).strftime('\n%b'), ha='center', va='center')

        for j, noise in enumerate(noises):

            color = ['#1f77b4', '#ff7f0e', '#2ca02c'][j]

            values = np.array([v[1][j] for v in data[month]]).flatten()
            plt.bar(x, values, width=1.0, color=color,
                    label=('$\sigma_{sat}$ = %g' % noise if month == 1 else '_none_'))

            text = '%d/%d' % (sum(values > 0), valid_days[month])

            if 'berlin' in filename:
                if tracer == 'CO2':
                    yy = [199, 192, 185][j]
                else:
                    yy = [695, 669][j]
            else:
                if tracer == 'CO2':
                    yy = [1190, 1150, 1110][j]
                else:
                    yy = [2190, 2130][j]
            plt.text(xx, yy, text, ha='center', va='top', color=color)



    plt.xticks(np.arange(14) * most_days - 0.5, [])
    plt.grid(True)
    plt.xlim(-0.5, 12 * most_days - 0.5)
    plt.ylim(0, ymax)
    plt.ylabel('number of pixels in plume')

    if tracer in ['CO2', 'NO2']:
        plt.legend(loc='center right')

    plt.tight_layout()

    return data, valid_days




def read_signal_and_background(time, tracer, ref, oth, input_files,
        instrument='', orbit=0, lon_eq=0):
    """\
    """
    if time >= datetime(2015,7,1):
        starttime = datetime(2015,7,1)
    else:
        starttime = datetime(2015,1,1)

    nc_filename = sc.io.create_filename(time, starttime, input_files,
        instrument, orbit, lon_eq, path=cfg.ROOT
    )

    if tracer == 'CO2':
        if ref == 'bv':
            b0,b1,b2, clct = sc.io.read_fields(nc_filename, ['XCO2_B0', 'XCO2_B1', 'XCO2_B2', 'CLCT'])

            if time.month in [1,7]:
                signal = 0.55 * b0 + b1 + b2
            else:
                signal = b0 + b1 + b2
        else:
            signal, clct = sc.io.read_fields(nc_filename, ['XCO2_%s' % ref, 'CLCT'])

        bg, anth, jv, gpp, ra = sc.io.read_fields(nc_filename,
                ['XCO2_BG', 'XCO2_A', 'XCO2_%s' % oth, 'XCO2_GPP', 'XCO2_RA']
        )
        background = bg + anth + jv - gpp + ra

    elif tracer == 'NO2':
        if ref == 'bv':
            b0,b1,b2, clct = sc.io.read_fields(nc_filename, ['NO2_B0', 'NO2_B1', 'NO2_B2', 'CLCT'])

            if time.month in [1,7]:
                signal = 0.55 * b0 + b1 + b2
            else:
                signal = b0 + b1 + b2
        else:
            signal, clct = sc.io.read_fields(nc_filename, ['NO2_%s' % ref, 'CLCT'])

        signal, clct = sc.io.read_fields(nc_filename, ['NO2_%s' % ref, 'CLCT'])
        bg, anth, jv = sc.io.read_fields(nc_filename, ['NO2_BG', 'NO2_A', 'NO2_%s' % oth])
        background = bg + anth + jv

    elif tracer == 'CO':
        signal, clct = sc.io.read_fields(nc_filename, ['CO_%s' % ref, 'CLCT'])
        bg, anth, jv = sc.io.read_fields(nc_filename, ['CO_BG', 'CO_A', 'CO_%s' % oth])
        background = bg + anth + jv

    # spatial binning
    if input_files != 'level2':
        signal = sc.misc.spatial_binning(signal)[0]
        clct = sc.misc.spatial_binning(clct)[0]
        background = sc.misc.spatial_binning(background)[0]

    return signal, clct, background




def main2(filename, tracer, start, stop, plume_name='Berlin',
        input_files='2D', instrument='Sentinel_7_CO2'):

    if plume_name == 'Berlin':
        ref = 'BV'
        oth = 'JV'
    else:
        ref = 'JV'
        oth = 'BV'

    if tracer == 'CO2':
        noises = [0.5, 0.7, 1.0]
    elif tracer == 'NO2':
        noises = [1e15, 2e15]
    elif tracer == 'CO':
        noises = [4e17]

    struct = ndi.generate_binary_structure(2, 2)

    with open(filename, 'w') as outfile:

        outfile.write('UTC')
        for noise in noises:
            outfile.write('\tall pixels %g' % noise)
            outfile.write('\tcloud-free pixels %g' % noise)
            outfile.write('\tbackground std %g' % noise)
        outfile.write('\n')


        for time in sc.misc.iter_times(start, stop, timedelta(hours=24)):
            print(time)

            signal, clct, background = read_signal_and_background(time, tracer, ref, oth, input_files,
                    instrument='', orbit=0, lon_eq=0)


            outfile.write(time.strftime('%Y-%m-%d %H:%M'))

            for noise in noises:

                # background area is a diluted core plume 
                core_plume = signal >= noise
                bg_area = ndi.binary_dilation(core_plume, struct)

                # std in background area
                bg_std = background[bg_area].std()

                # really detectable for signal >= level
                level = np.sqrt(noise**2 + bg_std**2)
                n_pixels = np.sum(signal >= level)
                n_cc_pixels = np.sum(signal[clct < 0.01] >= level)

                # write
                outfile.write('\t%d\t%d\t%e' % (n_pixels, n_cc_pixels, bg_std))

            outfile.write('\n')



def find_plume(signal, background, noise):

    while condition():
        pass


def main(filename, start, stop, mask_clouds):


    field_names = [
        'XCO2_BV',  # defines the Berlin plume

        'XCO2_BG',  # background signal
        'XCO2_A',   # anthropogenic signal (except large power plants)
        'XCO2_JV',  # anthropogenic signal from large power plants

        'XCO2_RA',  # biosphere
        'XCO2_GPP',

        'NO2_BV_column', # define Berlin plume
        'NO2_BG_column', # background
        'NO2_A_column',
        'NO2_JV_column',

        'CO_BV_column', # define Berlin plume
        'CO_BG_column', # background
        'CO_A_column',
        'CO_JV_column',
    ]

    levels = [0.1, 0.3, 0.5, 0.7, 0.9, 1.1]


    with open(filename, 'w') as outfile:

        outfile.write('utc')

        for name in field_names:
            for level in levels:
                outfile.write('\tnumber of pixels > %.1f' % (level))
                outfile.write('\t%s_mean>%.1f' % (name, level))
                outfile.write('\t%s_std>%.1f' % (name, level))
        outfile.write('\n')


        for time in sc.misc.iter_times(start, stop, timedelta(hours=24)):
            print(time)
            outfile.write('%s' % time.strftime('%Y-%m-%d %H:%M'))

            # starttime of simulation
            if time >= datetime(2015,7,1):
                starttime = datetime(2015,7,1)
            else:
                starttime = datetime(2015,1,1)

            try:
                bv, clct = sc.io.read_field(time, starttime, ['XCO2_BV', 'CLCT'], input_files='2D')
            except OSError:
                continue

            bv = sc.misc.spatial_binning(bv)
            clct = sc.misc.spatial_binning(clct)

            for name in field_names:
                field = sc.io.read_field(time, starttime, [name], input_files='2D')
                field = sc.misc.spatial_binning(field)

                for level in levels:
                    if mask_clouds:
                        plume = (bv > level) & (clct < 0.01)
                    else:
                        plume = (bv > level)

                    n_pixels = np.sum(plume)

                    outfile.write('\t%d' % n_pixels)

                    if n_pixels:
                        outfile.write('\t%e' % field[plume].mean())
                        outfile.write('\t%e' % field[plume].std())
                    else:
                        outfile.write('\tNaN')
                        outfile.write('\tNaN')

            outfile.write('\n')



if __name__ == '__main__':

    start = datetime(2015,1,1,11)
    stop = datetime(2015,12,31,11)


    for plume_name in ['Berlin']:#, 'plants']:
        #filename = 'ts_co2_plume_stats_%s_corrected.dat' % plume_name.lower()
        #main2(filename, 'CO2', start, stop, plume_name)

        filename = 'ts_no2_plume_stats_%s_corrected.dat' % plume_name.lower()
        main2(filename, 'NO2', start, stop, plume_name)

        #filename = 'ts_co_plume_stats_%s.dat' % plume_name.lower()
        #main2(filename, 'CO', start, stop, plume_name)




