#!/usr/bin/env python

# Copyright 2018 Earth Sciences Department, BSC-CNS
#
# This file is part of HERMESv3_GR.
#
# HERMESv3_GR is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# HERMESv3_GR is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with HERMESv3_GR. If not, see <http://www.gnu.org/licenses/>.


import os
import sys

# ============== README ======================
"""
downloading website: contact to hugo.deniervandergon@tno.nl or jeroen.kuenen@tno.nl
reference: https://www.atmos-chem-phys.net/14/10963/2014/
Besides citing HERMESv3_GR, users must also acknowledge the use of the corresponding emission inventories in their works
"""

# ============== CONFIGURATION PARAMETERS ======================
INPUT_PATH = '/esarchive/recon/ecmwf/cams_reg_apv51_ref22/original_files'
OUTPUT_PATH = '/esarchive/recon/ecmwf/cams_reg_apv51_ref22/yearly_mean'
INPUT_NAME = 'Ref2_v2_0_PM_components_<year>_DT.csv'
FACTOR = 1. / (365. * 24. * 3600. * 1000.)  # To pass from g/year to Kg/s
POLLUTANT_INFO = {'ec_dt': {'input_name': 'EC_fine', 'unit_factor': FACTOR},
                  'oc_dt': {'input_name': 'OC_fine', 'unit_factor': FACTOR},
                  'pm25_dt': {'input_name': ['EC_fine', 'OC_fine', 'SO4_fine', 'Na_fine', 'OthMin_fine'],
                              'unit_factor': FACTOR},
                  'so4_dt': {'input_name': 'SO4_fine', 'unit_factor': FACTOR},
                  'pm10_dt': {'input_name': ['EC_fine', 'EC_coarse', 'OC_fine', 'OC_coarse', 'SO4_fine', 'SO4_coarse',
                                             'Na_fine', 'Na_coarse', 'OthMin_fine', 'OthMin_coarse'],
                              'unit_factor': FACTOR}}
# list_years = [2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015]
LIST_YEARS = [2018]
# ==============================================================


def calculate_grid_definition(in_path):
    """
    Calculate the latitude and longitude coordinates of the cell.

    :param in_path: Path to the file that contains all the information.
    :type in_path: str

    :return: Latitudes array, Longitudes array, Latitude interval, Longitude interval.
    :rtype: numpy.array, numpy.array, float, float
    """
    import pandas as pd
    import numpy as np

    dataframe = pd.read_csv(in_path, sep=',')
    # Not needed to discriminate cause point sources are reported at the center grid cell
    # dataframe = dataframe[dataframe.SourceType != 'P']

    # Longitudes
    lons = np.sort(np.unique(dataframe['Lon']))
    lons_interval = lons[1:] - lons[:-1]
    print('Lon min: {0}; Lon max: {1}; Lon inc: {2}; Lon num: {3}'.format(
        dataframe['Lon'].min(), dataframe['Lon'].max(), lons_interval.min(), len(lons)))

    # Latitudes
    lats = np.sort(np.unique(dataframe['Lat']))
    lats_interval = lats[1:] - lats[:-1]
    print('Lat min: {0}; Lat max: {1}; Lat inc: {2}; Lat num: {3}'.format(
        dataframe['Lat'].min(), dataframe['Lat'].max(), lats_interval.min(), len(lats)))

    lats = np.arange(-90 + lats_interval.min() / 2, 90, lats_interval.min(), dtype=np.float64)
    lons = np.arange(-180 + lons_interval.min() / 2, 180, lons_interval.min(), dtype=np.float64)

    return lats, lons, lats_interval.min(), lons_interval.min()


def create_pollutant_empty_info(gnfr_id, len_c_lats, len_c_lons):
    """
    Crate an empty pollutant list.

    :param gnfr_id: ID of the
    :type gnfr_id: str

    :param len_c_lats: Number of elements on the latitude array
    :type len_c_lats: int

    :param len_c_lons: Number of elements on the longitude array
    :type len_c_lons: int

    :return: Pollutant info
    :rtype: dict
    """
    import numpy as np

    pollutant_info = {}
    for pollutant_orig_name, pollutant_oirg_info in POLLUTANT_INFO.items():
        pollutant_name = pollutant_orig_name.replace('<GNFR_Sector>', gnfr_id)
        pollutant_info[pollutant_name] = pollutant_oirg_info.copy()
        pollutant_info[pollutant_name]['name'] = pollutant_name
        pollutant_info[pollutant_name]['units'] = 'kg.m-2.s-1'
        # aux_dict['units'] = 'Mg.km-2.year-1'
        pollutant_info[pollutant_name]['data'] = np.zeros((len_c_lats, len_c_lons))

    return pollutant_info


def do_transformation(year):
    """
    Make al the process to transform the emissions of the current year.

    :param year: year to process.
    :type year: int

    :return: True when everything finish well.
    :rtype: Bool
    """
    from hermesv3_gr.tools.netcdf_tools import write_netcdf, get_grid_area
    from hermesv3_gr.tools.coordinates_tools import create_bounds
    from datetime import datetime
    import pandas as pd
    import numpy as np

    in_file = os.path.join(INPUT_PATH, INPUT_NAME.replace('<year>', str(year)))

    c_lats, c_lons, lat_interval, lon_interval = calculate_grid_definition(in_file)

    b_lats = create_bounds(c_lats, number_vertices=2)
    b_lons = create_bounds(c_lons, number_vertices=2)

    dataframe = pd.read_csv(in_file, sep=',')

    dataframe.loc[:, 'row_lat'] = np.array((dataframe['Lat'] - (-90 + lat_interval / 2)) / lat_interval, dtype=np.int32)
    dataframe.loc[:, 'col_lon'] = np.array((dataframe['Lon'] - (-180 + lon_interval / 2)) / lon_interval, dtype=np.int32)

    for gnfr_id, gnfr_data in dataframe.groupby('GNFR_Sector'):
        print('gnfr', gnfr_id)
        pollutant_info = create_pollutant_empty_info(gnfr_id, len(c_lats), len(c_lons))
        # Other mobile sources ignoring sea cells (shipping emissions)
        #if gnfr_id == 8:
        #    for sea in ['ATL', 'BAS', 'BLS', 'MED', 'NOS']:
        #        gnfr_data = gnfr_data[gnfr_data.ISO3 != sea]

        gnfr_data = gnfr_data.groupby(['row_lat', 'col_lon']).sum().reset_index()

        for poll_name, poll_info in pollutant_info.items():
            if isinstance(poll_info['input_name'], list):
                for poll_in_name in poll_info['input_name']:
                    poll_info['data'][gnfr_data.row_lat, gnfr_data.col_lon] += gnfr_data[poll_in_name]
            else:
                poll_info['data'][gnfr_data.row_lat, gnfr_data.col_lon] += gnfr_data[poll_info['input_name']]
            poll_info['data'] = poll_info['data'].reshape((1,) + poll_info['data'].shape)

            aux_output_path = os.path.join(OUTPUT_PATH, '{0}_gnfr_{1}'.format(poll_name, gnfr_id))
            if not os.path.exists(aux_output_path):
                os.makedirs(aux_output_path, exist_ok=True)
            aux_output_path = os.path.join(aux_output_path, '{0}_{1}.nc'.format(poll_name, year))
            write_netcdf(aux_output_path, c_lats, c_lons, [poll_info], date=datetime(year, month=1, day=1),
                         boundary_latitudes=b_lats, boundary_longitudes=b_lons)
            cell_area = get_grid_area(aux_output_path)

            poll_info['data'] = poll_info['data'] * poll_info['unit_factor']/cell_area

            write_netcdf(aux_output_path, c_lats, c_lons, [poll_info], date=datetime(year, month=1, day=1),
                         boundary_latitudes=b_lats, boundary_longitudes=b_lons, cell_area=cell_area,
                         global_attributes={
                             'comment': 'Re-writing done by Carles Tena (carles.tena@bsc.es) from the BSC-CNS ' +
                                        '(Barcelona Supercomputing Center)'})
    return True


if __name__ == '__main__':
    for y in LIST_YEARS:
        do_transformation(y)



