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

"""
Extract emissions and emission factors for Berlin and Jaenschwalde
"""

import os

from datetime import datetime, timedelta
from netCDF4 import Dataset

import numpy as np

import smartcarb as sc


PATH = '/project/d73/input/emissions'

# area of COSMO grid cell
AREA = 1235941.1 # in m2


def point_sources(start, end, out_filename):
    file_exists = os.path.exists(out_filename)
    mode = 'a' if file_exists else 'w'

    facilities = sorted(sc.config.POINT_SOURCES.keys())
    with open(out_filename, mode) as outfile:
        if not file_exists:
            outfile.write('\t'.join(['UTC time'] + ['%s (Mt/y)' % f for f in facilities]))
            outfile.write('\n')

        for time in sc.misc.iter_times(start, end, timedelta(hours=1)):
            filename = os.path.join(PATH, time.strftime('CO2_CO_NOX_Berlin-big_%Y%m%d%H.nc'))

            print(time)
            outfile.write(str(time))

            with Dataset(filename) as nc:

                rlon = nc.variables['rlon'][:]
                rlat = nc.variables['rlat'][:]

                rlon, rlat = np.meshgrid(rlon, rlat)
                lon, lat = sc.misc.rotpole_to_wgs84(rlon, rlat, sc.config.pollon, sc.config.pollat)

                var = nc.variables['NOX_JV_E']

                for name in facilities:

                    lon_o, lat_o = sc.config.POINT_SOURCES[name]
                    i,j = sc.misc.nearest_index(lon_o, lat_o, lon, lat)

                    v = np.sum(var[:,:,i-2:i+3,j-2:j+3]) * AREA # kg s-1
                    v *= 0.031536 # kg/s -> Mt/y

                    outfile.write('\t')
                    outfile.write('%3.6f' % v)

            outfile.write('\n')
            outfile.flush()







def main(start, end, out_filename):

    file_exists = os.path.exists(out_filename)
    mode = 'a' if file_exists else 'w'

    with open(out_filename, mode) as outfile:

        if not file_exists:
            outfile.write('\t'.join(
                ['UTC time', 'constant emissions', 'variable emissions', 'B0', 'B1', 'B2'])
            )
            outfile.write('\n')

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

            filename = os.path.join(PATH, time.strftime('CO2_CO_NOX_Berlin-big_%Y%m%d%H.nc'))

            with Dataset(filename) as nc:
                variable = np.sum(nc.variables['CO2_BV_E'][:]) * AREA # kg s-1
                constant = np.sum(nc.variables['CO2_BC_E'][:]) * AREA # kg s-1
                b0 = np.sum(nc.variables['CO2_B0_E'][:]) * AREA # kg s-1
                b1 = np.sum(nc.variables['CO2_B1_E'][:]) * AREA # kg s-1
                b2 = np.sum(nc.variables['CO2_B2_E'][:]) * AREA # kg s-1

            line = '\t'.join([str(time), '%e' % constant, '%e' % variable, '%e' % b0, '%e' % b1 , '%e' % b2])
            outfile.write(line)
            outfile.write('\n')



if __name__ == '__main__':
    start = datetime(2015,1,1)
    end = datetime(2016,1,1)

    #main(start, end, '/users/gkuhl/projects/smartcarb/ts/ts_berlin_emissions.dat')
    point_sources(start, end, '/users/gkuhl/projects/smartcarb/ts/ts_point_sources_NOX_emissions.dat')






