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

import os
import sys

from datetime import datetime
from netCDF4 import Dataset
import numpy as np
import scipy.stats as stats

import smartcarb as sc

import config



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



def percentile_factory(per):
    def percentile(a):
        return stats.scoreatpercentile(a, per=per)

    return percentile


def main(filename, start, stop, field_names, compute_functions, vmins, vmaxs,
        area='Berlin', input_files='3D'):

    if area == 'Berlin':
        indices = [sc.io.berlin_slice() for _ in field_names]

    elif area == 'Bremen-TCCON': #53.1, 8.85
        indices = [(0, slice(361-2, 361+3), slice(71-2,71+3)) for _ in field_names]

    else:
        indices = Ellipsis

    # write headers
    with open(filename, 'w') as out_file:

        out_file.write('time')
        for name, function in zip(field_names, compute_functions):
            out_file.write(
                '\t%s_%s' % (name, function.__name__)
            )
        out_file.write('\n')


        for time in sc.misc.iter_times(start, stop):
            print(time)

            # starttime of simulation
            if time >= datetime(2015,7,1):
                starttime = datetime(2015,7,1)
                path = config.path_2d_fields_07
            else:
                starttime = datetime(2015,1,1)
                path = config.path_2d_fields_01

            try:
                values = sc.io.read_field(time, starttime, field_names, indices, input_files,
                        path=path)
                values = [func(v[(vmin <= v) & (v < vmax)])
                        for v, func, vmin, vmax in zip(values, compute_functions, vmins, vmaxs)
                ]
            except OSError:
                values = [np.nan for _ in compute_functions]

            line = '%s\t%s\n' % (
                    time.strftime('%Y-%m-%d %H:%M'),
                    '\t'.join('%e' % v for v in values)
            )
            out_file.write(line)





if __name__ == '__main__':


    scenario = 'tracers_bremen'

    if scenario == 'tracers_berlin':
        area = 'Berlin'
        input_files = '2D'

        field_names = ['XCO2_TOT', 'XCO2_BG', 'XCO2_GPP', 'XCO2_RA',
                'NO2_TOT_column', 'NO2_BG_column',
                'CO_TOT_column', 'CO_BG_column'
        ]
        n = len(field_names)
        vmins = [0.0 for _ in range(n)]
        vmaxs = [np.inf for _ in range(n)]

        compute_functions = [np.mean for _ in range(n)]

        start = datetime(2015,1,1)
        stop = datetime(2016,1,1)

        filename = 'ts_tracers_berlin.dat'


    elif scenario == 'tracers_bremen':
        area = 'Bremen-TCCON'
        input_files = '2D'

        field_names = ['XCO2_A', 'XCO2_BV', 'XCO2_JV', 'XCO2_BG', 'YCO2_BG', 'XCO2_GPP', 'XCO2_RA', 'XCO2_SURF', 'XCO2_NPR']
        n = len(field_names)
        vmins = [0.0    for _ in range(n)]
        vmaxs = [1000.0 for _ in range(n)]

        compute_functions = [np.mean for _ in range(n)]

        start = datetime(2015,1,1)
        stop = datetime(2016,1,1)

        filename = 'ts_xco2_bremen_tccon.dat'


    elif scenario == 'meteo_berlin':

        area = 'Berlin'
        input_files = '3D'

        field_names = ['T_2M', 'TD_2M', 'U_10M', 'V_10M', 'CLCT', 'PS', 'W_SNOW']
        n = len(field_names)

        vmins = [0.0, 0.0, -np.inf, -np.inf, 0.0, 0.0, 0.0]
        vmaxs = [np.inf for _ in range(n)]

        compute_functions = [np.mean for _ in range(n)]

        start = datetime(2015,1,1)
        stop = datetime(2016,1,1)

        filename = 'ts_meteo_berlin.dat'



    main(filename, start, stop, field_names, compute_functions, vmins, vmaxs,
            area, input_files)



