#!/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 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 *
# import pyextrae.sequential as pyextrae

global full_time


class Hermes(object):
    """
    Interface class for HERMESv3.
    """
    def __init__(self, config, new_date=None, comm=None):
        from hermesv3_gr.modules.grids.grid import 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, 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 = Grid.select_grid(
            self.comm,
            self.options.domain_type, self.options.vertical_description,  self.options.output_timestep_num,
            self.options.auxiliary_files_path, self.options.inc_lat, self.options.inc_lon, self.options.lat_orig,
            self.options.lon_orig, self.options.n_lat, self.options.n_lon, self.options.centre_lat,
            self.options.centre_lon, self.options.west_boundary, self.options.south_boundary, self.options.inc_rlat,
            self.options.inc_rlon, self.options.lat_1, self.options.lat_2, self.options.lon_0, self.options.lat_0,
            self.options.nx, self.options.ny, self.options.inc_x, self.options.inc_y, self.options.x_0,
            self.options.y_0, self.options.lat_ts)

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

        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,
            compress=settings.compressed_netcdf,
            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_BU 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)
        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 = Hermes(Config()).main()
    while date is not None:
        date = Hermes(Config(), new_date=date).main()
    sys.exit(0)


if __name__ == '__main__':
    run()
