#!/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 traceback
import numpy as np
import pandas as pd
import re
from hermesv3_gr import __version__
from hermesv3_gr.config import ConfigGr
from hermesv3_gr.config import configure_logger, log_message, precision
from hermesv3_gr.config import get_time_stamp, record_time, finalize_time_log
# from hermesv3_d.config import get_memory_stamp, record_memory, finalize_mem_log
from hermesv3_gr.modules.grid import select_grid, add_local_dates
from hermesv3_gr.modules.emision_inventories import EmissionInventory, GfasEmissionInventory, GfasHourlyEmissionInventory, PointGfasEmissionInventory, PointGfasHourlyEmissionInventory, PointSourceEmissionInventory
from hermesv3_gr.modules.vertical.vertical import VerticalDistribution
from hermesv3_gr.modules.temporal.temporal import TemporalDistribution
from mpi4py import MPI
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from typing import Optional

global full_time


class HermesGr(object):
    """
    Interface class for HERMESv3_GR.

    Attributes
    ----------
    comm : MPI.COMMUNICATOR
    rank : int
        MPI rank
    master : bool
        Indicates if the process is the master one.
    size : int
        Total amount of processes involved.
    options : NameSpace
        Configuration options
    levels : List[float]
        Vertical level description
    grid : Nes
        Destination grid
    emission_list : List[EmissionInventory]
        Emission inventory list
    time_step_list : List[int]
        List of hours between time-steps. e.g. [0, 1, 2, ..., 24]
    """
    def __init__(self, config: ConfigGr, new_date: Optional[datetime] = None, comm: Optional[MPI.Comm] = None):
        """

        Parameters
        ----------
        config : ConfigGr
            Configuration instance
        new_date : datetime or None
            New date to simulate
        comm : MPI.Communicator
        """
        global full_time
        st_time = full_time = get_time_stamp()

        # MPI Initialization
        if comm is None:
            comm = MPI.COMM_WORLD
        self.comm = comm
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        self.master = self.rank == 0

        self.config = config
        self.options = config.options

        # Updating starting date
        if new_date is not None:
            if new_date > self.options.end_date:
                return
            else:
                self.options.start_date = new_date
        
        # Logger
        configure_logger(log_path=str(
            os.path.join(self.options.output_dir, 'logs',
                         os.path.basename(self.config.get_output_name(self.options.start_date)))).replace('.nc', "_<rank>.log"),
                         log_detail=self.options.log_level)

        self.first_time = self.config.options.first_time

        log_message(f"===== Starting HERMESv3_GR {__version__} initialization =====")

        if self.options.output_model in ['CMAQ', 'WRF_CHEM'] and self.options.domain_type == 'global':
            log_message('ERROR: Check the .err file to get more info.')
            raise AttributeError('ERROR: Global domain is not available for {0} output model.'.format(
                self.options.output_model))

        self.vertical_description = self.config.options.vertical_description
        self.levels = VerticalDistribution.get_vertical_output_profile(self.vertical_description)
        log_message("Levels: {0}".format(self.levels), level=3)

        self.grid = select_grid(self.comm, self.options)
        self.grid.set_levels({'data': self.levels, 'units': 'm', 'positive': 'up'})

        self.time_step_list = TemporalDistribution.calculate_date_array(
            self.options.start_date, self.options.output_timestep_type, self.options.output_timestep_num,
            self.options.output_timestep_freq)

        self.add_time_info_to_grid()

        self.emission_list = self.make_emission_list()

        self.out_format = self.options.output_model

        log_message('===== End of HERMESv3_GR initialization =====\n')
        record_time('Init', 'TOTAL', get_time_stamp() - st_time)

    def add_time_info_to_grid(self):
        st_time = get_time_stamp()
        shape_3d = (len(self.time_step_list), self.grid.lat['data'].shape[0], self.grid.lon['data'].shape[-1])
        # Empty 2D (time, lat*lon)
        self.grid.variables['month'] = {'data': np.empty((shape_3d[0], shape_3d[1] * shape_3d[2]), dtype=np.uint8)}
        self.grid.variables['weekday'] = {'data': np.empty((shape_3d[0], shape_3d[1] * shape_3d[2]), dtype=np.uint8)}
        self.grid.variables['hour'] = {'data': np.empty((shape_3d[0], shape_3d[1] * shape_3d[2]), dtype=np.uint8)}
        self.grid.variables['julian'] = {'data': np.empty((shape_3d[0], shape_3d[1] * shape_3d[2]), dtype=np.uint16)}

        for i_time, aux_date in enumerate(self.time_step_list):
            pandas_dates = add_local_dates(self.grid, aux_date, pandas=True)
            # Filling
            self.grid.variables['month']['data'][i_time] = pandas_dates.dt.month
            self.grid.variables['weekday']['data'][i_time] = pandas_dates.dt.dayofweek
            self.grid.variables['hour']['data'][i_time] = pandas_dates.dt.hour
            self.grid.variables['julian']['data'][i_time] = pandas_dates.dt.dayofyear
        # 2D (time, lat*lon) to 3D (time, lat, lon)
        self.grid.variables['month']['data'] = self.grid.variables['month']['data'].reshape(shape_3d)
        self.grid.variables['weekday']['data'] = self.grid.variables['weekday']['data'].reshape(shape_3d)
        self.grid.variables['hour']['data'] = self.grid.variables['hour']['data'].reshape(shape_3d)
        self.grid.variables['julian']['data'] = self.grid.variables['julian']['data'].reshape(shape_3d)

        record_time('HERMES', 'add_time_info_to_grid', get_time_stamp() - st_time)
        return None

    def make_emission_list(self):
        """
        Extract the information of the cross table to read all the needed emissions.

        Returns
        -------
        List[EmissionInventory]
            Emission inventory list
        """
        st_time = get_time_stamp()

        path = self.options.cross_table
        df = pd.read_csv(path, sep=';', index_col=False)
        for column in ['ei', 'sector', 'ref_year', 'active', 'factor_mask', 'regrid_mask', 'pollutants', 'path',
                       'frequency', 'source_type', 'coverage', 'p_vertical', 'p_month', 'p_week', 'p_day', 'p_hour',
                       'p_speciation']:
            df_cols = list(df.columns.values)
            if column not in df_cols:
                raise AttributeError('ERROR: Column {0} is not in the {1} file.'.format(column, path))

        df = df[df['active'] == 1]
        num = 1
        emission_inventory_list = []
        for i, src_emission_inventory in df.iterrows():
            log_message('Loading emission {0}/{1} (Inventory: {2}; Sector: {3})'.format(
                num, len(df), src_emission_inventory.ei, src_emission_inventory.sector), level=1)
            num += 1
            pollutants = list(map(str, re.split(', |,|; |;| ', src_emission_inventory.pollutants)))

            # MONTHLY
            try:
                # gridded temporal profile
                p_month = src_emission_inventory.p_month.replace('<data_path>', self.options.data_path)
            except AttributeError:
                p_month = src_emission_inventory.p_month
            # WEEKLY
            try:
                # gridded temporal profile
                p_week = src_emission_inventory.p_week.replace('<data_path>', self.options.data_path)
            except AttributeError:
                p_week = src_emission_inventory.p_week
            # DAILY
            try:
                # gridded temporal profile
                p_day = src_emission_inventory.p_day.replace('<data_path>', self.options.data_path)
            except AttributeError:
                p_day = src_emission_inventory.p_day
            # HOURLY
            try:
                # gridded temporal profile
                p_hour = src_emission_inventory.p_hour.replace('<data_path>', self.options.data_path)
            except AttributeError:
                p_hour = src_emission_inventory.p_hour

            emission_inventory_path = src_emission_inventory.path.replace('<data_path>', self.options.data_path)
            emission_inventory_path = emission_inventory_path.replace('<input_dir>', self.options.input_dir)

            if src_emission_inventory.source_type == 'area':
                if src_emission_inventory.ei[:4] == 'GFAS':
                    if src_emission_inventory.frequency == 'daily':
                        emission_inventory_list.append(
                            GfasEmissionInventory(
                                self.options, self.grid, self.options.start_date, src_emission_inventory.ei,
                                src_emission_inventory.source_type,
                                src_emission_inventory.sector, pollutants, emission_inventory_path,
                                src_emission_inventory.frequency, self.levels,
                                reference_year=src_emission_inventory.ref_year,
                                coverage=src_emission_inventory.coverage,
                                factors=src_emission_inventory.factor_mask,
                                regrid_mask=src_emission_inventory.regrid_mask,
                                p_vertical=src_emission_inventory.p_vertical, p_month=p_month, p_week=p_week, p_day=p_day,
                                p_hour=p_hour, p_speciation=src_emission_inventory.p_speciation,
                                countries_shapefile=self.options.countries_shapefile))
                    elif src_emission_inventory.frequency == 'hourly':
                        emission_inventory_list.append(
                            GfasHourlyEmissionInventory(
                                self.options, self.grid, self.options.start_date, src_emission_inventory.ei,
                                src_emission_inventory.source_type,
                                src_emission_inventory.sector, pollutants, emission_inventory_path,
                                src_emission_inventory.frequency, self.levels,
                                reference_year=src_emission_inventory.ref_year,
                                coverage=src_emission_inventory.coverage,
                                factors=src_emission_inventory.factor_mask,
                                regrid_mask=src_emission_inventory.regrid_mask,
                                p_vertical=src_emission_inventory.p_vertical, p_month=p_month, p_week=p_week,
                                p_day=p_day,
                                p_hour=p_hour, p_speciation=src_emission_inventory.p_speciation,
                                countries_shapefile=self.options.countries_shapefile))
                else:
                    emission_inventory_list.append(
                        EmissionInventory(
                            self.options, self.grid, self.options.start_date, src_emission_inventory.ei,
                            src_emission_inventory.source_type,
                            src_emission_inventory.sector, pollutants, emission_inventory_path,
                            src_emission_inventory.frequency, self.levels,
                            reference_year=src_emission_inventory.ref_year,
                            coverage=src_emission_inventory.coverage,
                            factors=src_emission_inventory.factor_mask,
                            regrid_mask=src_emission_inventory.regrid_mask,
                            p_vertical=src_emission_inventory.p_vertical, p_month=p_month, p_week=p_week, p_day=p_day,
                            p_hour=p_hour, p_speciation=src_emission_inventory.p_speciation,
                            countries_shapefile=self.options.countries_shapefile))
            elif src_emission_inventory.source_type == 'point':
                if src_emission_inventory.ei[:4] == 'GFAS':
                    if src_emission_inventory.frequency == 'daily':
                        emission_inventory_list.append(
                            PointGfasEmissionInventory(
                                self.options, self.grid, self.options.start_date, src_emission_inventory.ei,
                                src_emission_inventory.source_type,
                                src_emission_inventory.sector, pollutants, emission_inventory_path,
                                src_emission_inventory.frequency, self.levels,
                                reference_year=src_emission_inventory.ref_year,
                                factors=src_emission_inventory.factor_mask,
                                regrid_mask=src_emission_inventory.regrid_mask,
                                p_vertical=src_emission_inventory.p_vertical,
                                p_month=p_month, p_week=p_week, p_day=p_day, p_hour=p_hour,
                                p_speciation=src_emission_inventory.p_speciation,
                                countries_shapefile=self.options.countries_shapefile))
                    elif src_emission_inventory.frequency == 'hourly':
                        emission_inventory_list.append(
                            PointGfasHourlyEmissionInventory(
                                self.options, self.grid, self.options.start_date, src_emission_inventory.ei,
                                src_emission_inventory.source_type,
                                src_emission_inventory.sector, pollutants, emission_inventory_path,
                                src_emission_inventory.frequency, self.levels,
                                reference_year=src_emission_inventory.ref_year,
                                factors=src_emission_inventory.factor_mask,
                                regrid_mask=src_emission_inventory.regrid_mask,
                                p_vertical=src_emission_inventory.p_vertical,
                                p_speciation=src_emission_inventory.p_speciation,
                                countries_shapefile=self.options.countries_shapefile))
                    else:
                        raise NotImplementedError("ERROR: {0} frequency are not implemented, use 'daily' or 'hourly.")
                else:
                    emission_inventory_list.append(
                        PointSourceEmissionInventory(self.options, self.grid, self.options.start_date,
                                                     src_emission_inventory.ei,
                                                     src_emission_inventory.source_type, src_emission_inventory.sector,
                                                     pollutants, emission_inventory_path,
                                                     src_emission_inventory.frequency, self.levels,
                                                     reference_year=src_emission_inventory.ref_year,
                                                     factors=src_emission_inventory.factor_mask,
                                                     regrid_mask=src_emission_inventory.regrid_mask,
                                                     p_vertical=src_emission_inventory.p_vertical,
                                                     p_month=p_month,
                                                     p_week=p_week,
                                                     p_day=p_day,
                                                     p_hour=p_hour,
                                                     p_speciation=src_emission_inventory.p_speciation,
                                                     countries_shapefile=self.options.countries_shapefile))
            else:
                raise ValueError("ERROR: The emission inventory source type '{0}'".format(
                    src_emission_inventory.source_type) +
                                 " is not implemented. Use 'area' or 'point'")
            self.comm.Barrier()
        record_time('EmissionInventory', 'make_emission_list', get_time_stamp() - st_time)

        return emission_inventory_list

    def prepare_inventories(self):
        """
        Proces all the emission inventories to calculate the vertical distribution and speciation
        """
        st_time = get_time_stamp()
        for num, ei in enumerate(self.emission_list):
            log_message('Processing emission inventory {0} for the sector {1} ({2}/{3}):'.format(
                ei.inventory_name, ei.sector, num + 1, len(self.emission_list)))

            if ei.vertical is not None:
                log_message("Vertical distribution.", level=2)
                if ei.source_type == 'area':
                    ei.vertical_factors = ei.vertical.calculate_weights()
                    del ei.vertical

                elif ei.source_type == 'point':
                    # To avoid use point source as area source when is going to apply vertical factors while writing
                    ei.vertical = None
                else:
                    raise AttributeError('Unrecognized emission source type {0}'.format(ei.source_type))
            else:
                if ei.source_type == 'area':
                    ei.vertical_factors = np.zeros(len(self.levels), dtype=precision)
                    ei.vertical_factors[0] = 1
                elif ei.source_type == 'point':
                    ei.vertical_factors = np.ones(len(self.levels), dtype=precision)
                else:
                    raise AttributeError('Unrecognized emission source type {0}'.format(ei.source_type))
            if ei.speciation is not None:
                ei.emissions = ei.speciation.do_speciation(ei.emissions)
        record_time('EmissionInventory', 'prepare_inventories', get_time_stamp() - st_time)
        return None

    def get_time_step_emissions(self, time_step, i_time):
        """
        Obtain the emissions of the selected time step

        Parameters
        ----------
        time_step : datetime
            Time stamp to simulate
        i_time : int
            Index of the time-step

        Returns
        -------
        dict
            Time stamp emissions in Nes format
        """
        st_time = get_time_stamp()
        result = None
        for num, ei in enumerate(self.emission_list):
            if ei.temporal is not None:
                ei.temporal_factors = np.array([ei.temporal.calculate_2d_temporal_factors(i_time)],
                                               dtype=precision)
            else:
                ei.temporal_factors = None

            if isinstance(ei, GfasHourlyEmissionInventory) or isinstance(ei, PointGfasHourlyEmissionInventory):
                # Needed to load hourly emissions
                ei.date = time_step

            if result is None:
                result = ei.get_4d_emissions()
            else:
                for var_name, var_info in ei.get_4d_emissions().items():
                    result[var_name]['data'] += var_info['data']

        record_time('EmissionInventory', 'get_time_step_emissions', get_time_stamp() - st_time)
        return result

    def prepare_output(self):
        """
        Prepares the output to be written

        Returns
        -------
        Nes
            Output Nes object with lazy variables
        """
        result = self.grid.copy(copy_vars=False)
        result.cell_measures = self.grid.cell_measures
        result.set_strlen(None)
        result.set_communicator(self.comm)
        result.set_levels({'data': np.array(self.levels, dtype=np.float32), 'units': 'm', 'positive': 'up'})
        result.set_time(self.time_step_list)
        result.variables = self.emission_list[0].prepare_lazy_output_emissions()

        # Global attributes
        result.global_attrs = {}
        if self.out_format == 'CMAQ':
            if self.options.output_attributes is not None:
                global_attributes = pd.read_csv(self.options.output_attributes, sep=',', index_col='attribute')
                for attr_name in ['EXEC_ID', 'FTYPE', 'NTHIK', 'VGTYP', 'VGTOP', 'VGLVLS', 'GDNAM']:
                    if attr_name in global_attributes.index:
                        result.global_attrs[attr_name] = global_attributes.loc[attr_name, 'value']

            result.global_attrs['FILEDESC'] = 'Emissions generated by HERMESv3_GR.'
        elif self.out_format == 'WRF_CHEM':
            if self.options.output_attributes is not None:
                global_attributes = pd.read_csv(self.options.output_attributes, sep=',', index_col='attribute')
                for attr_name in [
                    'DAMPCOEF', 'KHDIF', 'KVDIF', 'CEN_LAT', 'CEN_LON', 'DT', 'BOTTOM-TOP_GRID_DIMENSION', 'DIFF_OPT',
                    'KM_OPT', 'DAMP_OPT', 'MP_PHYSICS', 'RA_LW_PHYSICS', 'RA_SW_PHYSICS', 'SF_SFCLAY_PHYSICS',
                    'SF_SURFACE_PHYSICS', 'BL_PBL_PHYSICS', 'CU_PHYSICS', 'SF_LAKE_PHYSICS', 'SURFACE_INPUT_SOURCE',
                    'SST_UPDATE', 'GRID_FDDA', 'GFDDA_INTERVAL_M', 'GFDDA_END_H', 'GRID_SFDDA', 'SGFDDA_INTERVAL_M',
                    'SGFDDA_END_H', 'BOTTOM-TOP_PATCH_START_UNSTAG', 'BOTTOM-TOP_PATCH_END_UNSTAG',
                    'BOTTOM-TOP_PATCH_START_STAG', 'BOTTOM-TOP_PATCH_END_STAG', 'GRID_ID', 'PARENT_ID',
                    'I_PARENT_START', 'J_PARENT_START', 'PARENT_GRID_RATIO', 'NUM_LAND_CAT', 'ISWATER', 'ISLAKE',
                    'ISICE', 'ISURBAN', 'ISOILWATER', 'GRIDTYPE', 'MMINLU']:
                    if attr_name in global_attributes.index:
                        result.global_attrs[attr_name] = global_attributes.loc[attr_name, 'value']
        result.global_attrs['HISTORY'] = \
            "Code developed by Barcelona Supercomputing Center (BSC, https://www.bsc.es/). " + \
            "Developer: Carles Tena Medina (carles.tena@bsc.es) " + \
            "Reference: Guevara et al., 2019, GMD. https://doi.org/10.5194/gmd-12-1885-2019"
        return result

    def main(self):
        """
        Main functionality of the model.

        Returns
        -------
        datetime
            New datetime if other simulation is requested, otherwise None
        """
        if self.first_time:
            # Stop run
            log_message('===== HERMESv3_GR First Time finished successfully =====')

        else:
            write_time = 0
            st_time = get_time_stamp()
            log_message('')
            log_message('===== Starting HERMESv3_GR Calculation =====')

            self.prepare_inventories()  # To obtain speciated emissions and vertical factors

            # Preparing output file
            aux_time = get_time_stamp()
            result = self.prepare_output()
            out_path = self.config.get_output_name(self.options.start_date)
            result.to_netcdf(out_path, serial=self.options.serial_write, keep_open=True, nc_type=self.out_format)
            write_time += get_time_stamp() - aux_time

            log_message(f'Created empty output file {out_path}', level=1)

            for i_time, time_step in enumerate(self.time_step_list):
                log_message('Calculating {0} data ({1}/{2})'.format(
                    time_step, i_time+1, len(self.time_step_list)), level=2)
                result.variables = self.get_time_step_emissions(time_step, i_time)
                log_message('Writing data', level=3)
                aux_time = get_time_stamp()
                result.append_time_step_data(i_time, out_format=self.out_format)
                write_time += get_time_stamp() - aux_time

            log_message(f"***** Output file done: {out_path}")

            # Closing output file
            aux_time = get_time_stamp()
            result.close()
            write_time += get_time_stamp() - aux_time

            log_message("===== HERMESv3_GR simulation finished successfully =====")
            record_time('HERMES', 'Main', get_time_stamp() - st_time)
            record_time('HERMES', 'Write', write_time)
        record_time('HERMES', 'TOTAL', get_time_stamp() - full_time)

        finalize_time_log(
            str(os.path.join(self.options.output_dir, 'logs', os.path.basename(
                self.config.get_output_name(self.options.start_date)))).replace(
                '.nc', f"_Times_{str(self.size).zfill(4)}.csv"))

        # Updating the date for a date-loop HERMES simulation
        if self.options.start_date < self.options.end_date:
            if self.options.output_timestep_type in ['hourly', 'daily']:
                new_date = self.options.start_date + timedelta(days=1)
            elif self.options.output_timestep_type in ['monthly']:
                new_date = self.options.start_date + relativedelta(months=1)
            elif self.options.output_timestep_type in ['yearly']:
                new_date = self.options.start_date + relativedelta(years=1)
            else:
                raise RuntimeError("Unknown timestep_type")
            return new_date

        return None


def run():
    try:
        date = HermesGr(ConfigGr()).main()
        while date is not None:
            date = HermesGr(ConfigGr(), new_date=date).main()
    except Exception as e:
        log_message('CRITICAL ERROR {e}')
        log_message(f"Process {MPI.COMM_WORLD.Get_rank()}: Critical error detected {e}, aborting MPI.")
        log_message(f"Process {MPI.COMM_WORLD.Get_rank()}: Traceback:\n{traceback.format_exc()}", level=9)

        MPI.COMM_WORLD.Abort(1)


if __name__ == '__main__':
    run()
