#!/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 timeit
from traceback import format_exception

from hermesv3_gr.config import settings
from hermesv3_gr.config.config import Config
from hermesv3_gr.modules.emision_inventories.emission_inventory import EmissionInventory
from hermesv3_gr.modules.vertical.vertical import VerticalDistribution
from hermesv3_gr.tools.netcdf_tools import *

global full_time


class HermesGr(object):
    """
    Interface class for HERMESv3_GR.
    """
    def __init__(self, config, new_date=None, comm=None):
        from hermesv3_gr.modules.grids.grid import select_grid
        from hermesv3_gr.modules.temporal.temporal import TemporalDistribution
        from hermesv3_gr.modules.writing.writer import Writer
        global full_time
        st_time = full_time = timeit.default_timer()

        if comm is None:
            comm = MPI.COMM_WORLD
        self.comm = comm
        self.config = config
        self.options = config.options

        # Updating starting date
        if new_date is not None:
            self.options.start_date = new_date

        config.set_log_level(self.comm)
        settings.define_log_file(self.options.output_dir, config.get_output_name(self.options.start_date))
        settings.define_times_file()

        settings.write_log('Starting HERMESv3 initialization:')

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

        self.levels = VerticalDistribution.get_vertical_output_profile(self.options.vertical_description)

        self.grid = select_grid(self.comm, self.options)

        self.emission_list = EmissionInventory.make_emission_list(self.comm, self.options, self.grid, self.levels,
                                                                  self.options.start_date,
                                                                  self.options.countries_shapefile)

        self.delta_hours = TemporalDistribution.calculate_delta_hours(
            self.options.start_date, self.options.output_timestep_type, self.options.output_timestep_num,
            self.options.output_timestep_freq)

        self.writer = Writer.get_writer(
            self.options.output_model, self.config.get_output_name(self.options.start_date), self.grid,
            self.levels, self.options.start_date, self.delta_hours, self.options.output_attributes,
            self.options.compression_level, parallel=not settings.writing_serial)

        settings.write_log('End of HERMESv3 initialization.')
        settings.write_time('HERMES', 'Init', timeit.default_timer() - st_time, level=1)

    # @profile
    def main(self, return_emis=False):
        """
        Main functionality of the model.
        """
        from datetime import timedelta

        if self.options.first_time:
            # Stop run
            settings.write_log('***** HERMESv3_GR First Time finished successfully *****')
        else:
            st_time = timeit.default_timer()
            settings.write_log('')
            settings.write_log('***** Starting HERMESv3_GR *****')
            num = 1
            for ei in self.emission_list:
                settings.write_log('Processing emission inventory {0} for the sector {1} ({2}/{3}):'.format(
                    ei.inventory_name, ei.sector, num, len(self.emission_list)))
                num += 1

                ei.do_regrid()

                if ei.vertical is not None:
                    settings.write_log("\tCalculating vertical distribution.", level=2)
                    if ei.source_type == 'area':
                        ei.vertical_factors = ei.vertical.calculate_weights()
                    elif ei.source_type == 'point':
                        ei.calculate_altitudes(self.options.vertical_description)
                        ei.point_source_by_cell()
                        # To avoid use point source as area source when is going to apply vertical factors while writing
                        ei.vertical = None
                    else:
                        settings.write_log('ERROR: Check the .err file to get more info.')
                        if settings.rank == 0:
                            raise AttributeError('Unrecognized emission source type {0}'.format(ei.source_type))
                        sys.exit(1)

                if ei.temporal is not None:
                    ei.temporal_factors = ei.temporal.calculate_3d_temporal_factors()
                if ei.speciation is not None:
                    ei.emissions = ei.speciation.do_speciation(ei.emissions)
            if return_emis:
                return self.writer.get_emis(self.emission_list)
            else:
                self.writer.write(self.emission_list)

            settings.write_log("***** HERMESv3_GR simulation finished successfully *****")
            settings.write_time('HERMES', 'main', timeit.default_timer() - st_time)
        settings.write_time('HERMES', 'TOTAL', timeit.default_timer() - full_time)

        if settings.log_level >= 3:
            settings.finish_logs(self.options.output_dir, self.options.start_date)

        if self.options.start_date < self.options.end_date:
            return self.options.start_date + timedelta(days=1)

        return None


def run():
    date = HermesGr(Config()).main()
    while date is not None:
        date = HermesGr(Config(), new_date=date).main()
    sys.exit(0)


if __name__ == '__main__':
    run()

# def mpiabort_excepthook(error_type, error_value, error_traceback):
#     """
#     Override sys.excepthookand call explicitly MPI.COMM_WORLD.Abort in it.
#
#     https://stackoverflow.com/questions/49868333/fail-fast-with-mpi4py
#     """
#     msg = "{orange}Rank {rank:03d} has raised a {end_c}{red}{type}{end_c}: {value}\n".format(
#         red='\033[91m', orange='\033[93m', end_c='\033[0m', value=error_value,
#         rank=MPI.COMM_WORLD.Get_rank(), type=str(error_type).replace("<class '", "").replace("'>", ""))
#     msg += ''.join(format_exception(error_type, error_value, error_traceback))
#
#     print(msg)
#     MPI.COMM_WORLD.Abort()
#     # noinspection PyUnreachableCode
#     sys.__excepthook__(error_type, error_value, error_traceback)
#
#
# if __name__ == '__main__':
#     sys.excepthook = mpiabort_excepthook
#     run()
#     sys.excepthook = sys.__excepthook__

