#!/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
import numpy as np
import pandas as pd
import timeit
from datetime import datetime, timedelta
import xarray as xr

# ============== README ======================
"""
downloading website: 
reference: 
Besides citing HERMESv3_GR, users must also acknowledge the use of the corresponding emission inventories in their works
"""

# ============== CONFIGURATION PARAMETERS ======================
# The original emissions are downloaded from the following link:
# downloading website: http://geoschemdata.wustl.edu/ExtData/HEMCO/OFFLINE_SOILNOX/v2021-12/0.25x0.3125/
# reference: https://doi.org/10.1038/s41597-020-0488-5

INPUT_PATH = '/esarchive/recon/pku/hemco_soilnox-v202112/original_files/<MM>/soilnox_025.<YYYYMMDD>.nc'
OUTPUT_PATH = '/esarchive/recon/pku/hemco_soilnox-v202112/daily_mean/<pollutant>_<sector>/<pollutant>_<YYYYMMDD>.nc'
SECTOR_LIST = ['soil']
POLLUTANT_INFO = {"EmisNO_Fert": "nox_n_fert", "SOIL_NOx": "nox_n_tot"}

START_DAY = datetime(year=2019, month=1, day=1)
END_DAY = datetime(year=2019, month=12, day=31)
# END_DAY = datetime(year=2019, month=1, day=1)

XARRAY = True
# ==============================================================


def create_bounds(coords, number_vertices=2):
    """
    Calculate the vertices coordinates.

    :param coords: Coordinates in degrees (latitude or longitude)
    :type coords: numpy.array

    :param number_vertices: Non mandatory parameter that informs the number of vertices that must have the boundaries.
            (by default 2)
    :type number_vertices: int

    :return: Array with as many elements as vertices for each value of coords.
    :rtype: numpy.array
    """
    import numpy as np

    interval = coords[1] - coords[0]

    coords_left = coords - interval / 2
    coords_right = coords + interval / 2
    if number_vertices == 2:
        bound_coords = np.dstack((coords_left, coords_right))
    elif number_vertices == 4:
        bound_coords = np.dstack((coords_left, coords_right, coords_right, coords_left))
    else:
        raise ValueError('The number of vertices of the boudaries must be 2 or 4')

    return bound_coords


def parse_path(file_path, date_aux=None, sector=None, pollutant=None):
    """
    :param file_path: Path to parse
    :type file_path: str

    :param date_aux: Date to parse
    :type date_aux: datetime, None

    :param sector: Sector name
    :type sector: str, None

    :param pollutant: Pollutant name
    :type pollutant: str, None

    :return: Complete path
    :rtype: str
    """
    if date_aux is not None:
        file_path = file_path.replace('<YYYYMMDD>', date_aux.strftime("%Y%m%d"))
        file_path = file_path.replace('<MM>', date_aux.strftime("%m"))
    if pollutant is not None:
        file_path = file_path.replace('<pollutant>', pollutant)
    if sector is not None:
        file_path = file_path.replace('<sector>', sector)
    return file_path


def make_dir(file_path):
    parent_dir = os.path.dirname(file_path)
    if not os.path.exists(parent_dir):
        os.makedirs(parent_dir, exist_ok=True)
    return True


def make_grid(template_path):
    grid = xr.open_dataset(template_path)

    grid.coords['cell_area'] = grid['AREA']
    grid.coords['cell_area'].attrs = {"long_name": "area of the grid cell",
                                      "standard_name": "cell_area",
                                      "units": "m2"}

    grid.coords['crs'] = int(0)
    grid.coords['crs'].attrs = {"grid_mapping_name": "latitude_longitude",
                                "semi_major_axis": 6371000.,
                                "inverse_flattening": 0}

    lat_bnds = create_bounds(grid.coords['lat'])
    lat_bnds = lat_bnds.reshape((lat_bnds.shape[-2], lat_bnds.shape[-1]))
    grid.coords['lat_bnds'] = (('lat', 'nv'), lat_bnds)
    grid.coords['lat'].attrs['bounds'] = 'lat_bnds'

    lon_bnds = create_bounds(grid.coords['lon'])
    lon_bnds = lon_bnds.reshape((lon_bnds.shape[-2], lon_bnds.shape[-1]))
    grid.coords['lon_bnds'] = (('lon', 'nv'), lon_bnds)
    grid.coords['lon'].attrs['bounds'] = 'lon_bnds'

    var_list = list(grid.keys())
    for var_name in var_list:
        del grid[var_name]

    return grid


def main_xarray():

    grid_nc = make_grid(parse_path(INPUT_PATH, date_aux=START_DAY))

    day_aux = START_DAY
    for sector in SECTOR_LIST:
        while day_aux <= END_DAY:
            print(day_aux)
            in_nc = xr.open_dataset(parse_path(INPUT_PATH, date_aux=day_aux))
            for poll_in_name, poll_out_name in POLLUTANT_INFO.items():
                out_nc = grid_nc.copy()
                out_nc[poll_out_name] = (('lat', 'lon'), (in_nc[poll_in_name][:].mean(dim='time').data))
                out_nc[poll_out_name].attrs['cell_measures'] = "area: cell_area"
                out_nc[poll_out_name].attrs['grid_mapping'] = "crs"

                out_nc['time'] = [day_aux]

                out_nc.attrs['comment'] = "Re-writing done by Carles Tena (carles.tena@bsc.es) from the BSC-CNS " + \
                                          "(Barcelona Supercomputing Center)"

                out_path = parse_path(OUTPUT_PATH, date_aux=day_aux, sector=sector, pollutant=poll_out_name)
                make_dir(out_path)
                out_nc.to_netcdf(out_path)

            day_aux += timedelta(days=1)

    return True


def main():
    day_aux = START_DAY

    while day_aux <= END_DAY:
        in_path = parse_path(INPUT_PATH, day_aux)

        day_aux += timedelta(days=1)

    return True


if __name__ == '__main__':
    initime = timeit.default_timer()

    if XARRAY:
        main_xarray()
    else:
        main()

    print("\t{0:.3f}s".format(timeit.default_timer() - initime))
    sys.exit(0)
