#!/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 MercatorGrid(Grid):
    """
    Mercator grid object that contains all the information to do a mercator 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 lon_0: Value of the Lon0 for the LCC grid type.
    :type lon_0: float

    :param nx: Number of cells on the x dimension.
    :type nx: int

    :param ny: Number of cells on the y dimension.
    :type ny: int

    :param inc_x: Increment between x dimensions cell centroids (metres).
    :type inc_x: int

    :param inc_y: Increment between y dimensions cell centroids (metres).
    :type inc_y: int

    :param x_0: Value of the X0 for the LCC grid type.
    :type x_0: float

    :param y_0: Value of the Y0 for the LCC grid type.
    :type y_0: float

    :param earth_radius: Radius of the Earth (metres).
    Default = 6370000.000
    :type earth_radius: float
    """

    def __init__(self, grid_type, vertical_description_path, timestep_num, temporal_path, lat_ts, lon_0,
                 nx, ny, inc_x, inc_y, x_0, y_0, earth_radius=6370000.000, comm=None):
        st_time = timeit.default_timer()
        settings.write_log('\tCreating Mercator grid.', level=2)

        # Initialises with parent class
        super(MercatorGrid, self).__init__(grid_type, vertical_description_path, temporal_path, comm)

        # Setting parameters
        self.lat_ts = lat_ts
        self.lon_0 = lon_0
        self.nx = nx
        self.ny = ny
        self.inc_x = inc_x
        self.inc_y = inc_y
        self.x_0 = x_0 + (inc_x / 2)
        self.y_0 = y_0 + (inc_y / 2)
        self.earth_radius = earth_radius

        # UTM coordinates
        self.x = None
        self.y = None

        # Creating coordinates
        self.crs = "+proj=merc +a={2} +b={2} +lat_ts={0} +lon_0={1}".format(self.lat_ts, self.lon_0, earth_radius)

        self.create_coords()

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

        self.calculate_bounds()

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

        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('MercatorGrid', 'Init', timeit.default_timer() - st_time, level=1)

    def write_coords_netcdf(self):
        """
        Writes the temporal file with the coordinates of the output needed to generate the weight matrix.
        If it is already well created it will only add the cell_area parameter.
        """
        from hermesv3_gr.tools.netcdf_tools import write_netcdf

        st_time = timeit.default_timer()

        if not self.chech_coords_file():
            # Writes an auxiliary empty NetCDF only with the coordinates and an empty variable.
            write_netcdf(self.coords_netcdf_file, self.center_latitudes, self.center_longitudes,
                         [{'name': 'var_aux', 'units': '', 'data': 0}],
                         boundary_latitudes=self.boundary_latitudes, boundary_longitudes=self.boundary_longitudes,
                         mercator=True, lcc_x=self.x, lcc_y=self.y, lon_0=self.lon_0, lat_ts=self.lat_ts)

            # Calculate the cell area of the auxiliary NetCDF file
            self.cell_area = self.get_cell_area()

            # Re-writes the NetCDF adding the cell area
            write_netcdf(self.coords_netcdf_file, self.center_latitudes, self.center_longitudes,
                         [
                             {'name': 'var_aux',
                              'units': '',
                              'data': 0}
                         ],
                         boundary_latitudes=self.boundary_latitudes, boundary_longitudes=self.boundary_longitudes,
                         cell_area=self.cell_area,
                         mercator=True, lcc_x=self.x, lcc_y=self.y, lon_0=self.lon_0, lat_ts=self.lat_ts)
        else:
            self.cell_area = self.get_cell_area()

        settings.write_time('MercatorGrid', 'write_coords_netcdf', timeit.default_timer() - st_time, level=3)

    def create_coords(self):
        """
        Create the coordinates for a lambert conformal conic domain.
        """
        import numpy as np
        from pyproj import Proj

        st_time = timeit.default_timer()

        # Create a regular grid in metres (Two 1D arrays)
        self.x = np.linspace(self.x_0, self.x_0 + (self.inc_x * (self.nx - 1)), self.nx, dtype=np.float)
        if len(self.x)//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.x)//2))
            sys.exit(1)
        self.y = np.linspace(self.y_0, self.y_0 + (self.inc_y * (self.ny - 1)), self.ny, dtype=np.float)

        # 1D to 2D
        x = np.array([self.x] * len(self.y))
        y = np.array([self.y] * len(self.x)).T

        # Create UTM bounds
        y_b = super(MercatorGrid, self).create_bounds(y, self.inc_y, number_vertices=4, inverse=True)
        x_b = super(MercatorGrid, self).create_bounds(x, self.inc_x, number_vertices=4)

        # Create the LCC projection
        projection = Proj(self.crs)

        # UTM to Mercator
        self.center_longitudes, self.center_latitudes = projection(x, y, inverse=True)
        self.boundary_longitudes, self.boundary_latitudes = projection(x_b, y_b, inverse=True)

        settings.write_time('MercatorGrid', 'create_coords', timeit.default_timer() - st_time, level=3)


if __name__ == '__main__':
    pass
