#!/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 timeit

import hermesv3_gr.config.settings as settings
from grid import Grid


class LatLonGrid(Grid):
    """
    Regional regular lat-lon grid object that contains all the information to do a global output.

    :param grid_type: Type of the output grid [global, rotated, lcc, mercator].
    :type grid_type: str

    :param vertical_description_path: Path to the file that contains the vertical description.
    :type vertical_description_path: str

    :param timestep_num: Number of timesteps.
    :type timestep_num: int

    :param temporal_path: Path to the temporal folder.
    :type temporal_path: str

    :param inc_lat: Increment between latitude centroids.
    :type inc_lat: float

    :param inc_lon: Increment between longitude centroids.
    :type inc_lon: float

    :param lat_orig: Location of the latitude of the corner of the first cell (down left).
    :type lat_orig: float

    :param lon_orig: Location of the longitude of the corner of the first cell (down left).
    :type lon_orig: float

    :param n_lat: Number of cells on the latitude direction.
    :type n_lat = int

    :param n_lon: Number of cells on the latitude direction.
    :type n_lon = int

    """

    def __init__(self, grid_type, vertical_description_path, timestep_num, temporal_path, inc_lat, inc_lon, lat_orig,
                 lon_orig, n_lat, n_lon):
        import ESMF

        st_time = timeit.default_timer()
        settings.write_log('\tCreating Regional regular lat-lon grid.', level=2)

        # Initialize the class using parent
        super(LatLonGrid, self).__init__(grid_type, vertical_description_path, temporal_path)

        self.inc_lat = inc_lat
        self.inc_lon = inc_lon
        self.lat_orig = lat_orig
        self.lon_orig = lon_orig
        self.n_lat = n_lat
        self.n_lon = n_lon

        self.crs = {'init': 'epsg:4326'}
        self.create_coords()

        if not os.path.exists(self.coords_netcdf_file):
            if settings.rank == 0:
                self.write_coords_netcdf()
            settings.comm.Barrier()

        self.esmf_grid = self.create_esmf_grid_from_file(self.coords_netcdf_file, sphere=False)

        self.x_lower_bound = self.esmf_grid.lower_bounds[ESMF.StaggerLoc.CENTER][1]
        self.x_upper_bound = self.esmf_grid.upper_bounds[ESMF.StaggerLoc.CENTER][1]
        self.y_lower_bound = self.esmf_grid.lower_bounds[ESMF.StaggerLoc.CENTER][0]
        self.y_upper_bound = self.esmf_grid.upper_bounds[ESMF.StaggerLoc.CENTER][0]

        self.shape = (timestep_num, len(self.vertical_description), self.x_upper_bound-self.x_lower_bound,
                      self.y_upper_bound-self.y_lower_bound)
        settings.comm.Barrier()

        total_area = self.get_cell_area()
        self.full_shape = (timestep_num, len(self.vertical_description), total_area.shape[-2], total_area.shape[-1])
        self.cell_area = total_area[self.x_lower_bound:self.x_upper_bound, self.y_lower_bound:self.y_upper_bound]

        settings.write_time('LatLonGrid', 'Init', timeit.default_timer() - st_time, level=1)

    def create_coords(self):
        """
        Create the coordinates for a global domain.
        """
        import numpy as np

        st_time = timeit.default_timer()
        settings.write_log('\t\tCreating global coordinates', level=3)

        # From corner latitude /longitude to center ones
        lat_c_orig = self.lat_orig + (self.inc_lat / 2)
        lon_c_orig = self.lon_orig + (self.inc_lon / 2)

        self.center_latitudes = np.linspace(lat_c_orig, lat_c_orig + (self.inc_lat * (self.n_lat - 1)), self.n_lat,
                                            dtype=np.float)
        boundary_latitudes = self.create_bounds(self.center_latitudes, self.inc_lat)

        # ===== Longitudes =====
        self.center_longitudes = np.linspace(lon_c_orig, lon_c_orig + (self.inc_lon * (self.n_lon - 1)), self.n_lon,
                                             dtype=np.float)
        if len(self.center_longitudes)//2 < settings.size:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError("ERROR: Maximum number of processors exceeded. " +
                                     "It has to be less or equal than {0}.".format(len(self.center_longitudes)//2))
            sys.exit(1)

        boundary_longitudes = self.create_bounds(self.center_longitudes, self.inc_lon)

        self.boundary_latitudes = boundary_latitudes.reshape((1,) + boundary_latitudes.shape)
        self.boundary_longitudes = boundary_longitudes.reshape((1,) + boundary_longitudes.shape)

        settings.write_time('LatLonGrid', 'create_coords', timeit.default_timer() - st_time, level=2)


if __name__ == '__main__':
    pass
