#!/usr/bin/env python

import numpy as np
from copy import deepcopy


def calculate_grid_area(self):
    """
    Get coordinate bounds and call function to calculate the area of each cell of a grid.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object.
    """
        
    # Create bounds if they do not exist
    if self._lat_bnds is None or self._lon_bnds is None:
        self.create_spatial_bounds()

    # Get spatial number of vertices
    spatial_nv = self.lat_bnds['data'].shape[-1]

    # Reshape bounds
    if spatial_nv == 2:

        aux_shape = (self.lat_bnds['data'].shape[0], self.lon_bnds['data'].shape[0], 4)
        lon_bnds_aux = np.empty(aux_shape)
        lon_bnds_aux[:, :, 0] = self.lon_bnds['data'][np.newaxis, :, 0]
        lon_bnds_aux[:, :, 1] = self.lon_bnds['data'][np.newaxis, :, 1]
        lon_bnds_aux[:, :, 2] = self.lon_bnds['data'][np.newaxis, :, 1]
        lon_bnds_aux[:, :, 3] = self.lon_bnds['data'][np.newaxis, :, 0]

        lon_bnds = lon_bnds_aux
        del lon_bnds_aux

        lat_bnds_aux = np.empty(aux_shape)
        lat_bnds_aux[:, :, 0] = self.lat_bnds['data'][:, np.newaxis, 0]
        lat_bnds_aux[:, :, 1] = self.lat_bnds['data'][:, np.newaxis, 0]
        lat_bnds_aux[:, :, 2] = self.lat_bnds['data'][:, np.newaxis, 1]
        lat_bnds_aux[:, :, 3] = self.lat_bnds['data'][:, np.newaxis, 1]

        lat_bnds = lat_bnds_aux
        del lat_bnds_aux

    else:
        lon_bnds = self.lon_bnds['data']
        lat_bnds = self.lat_bnds['data']

    # Reshape bounds and assign as grid corner coordinates
    grid_corner_lon = deepcopy(lon_bnds).reshape(lon_bnds.shape[0]*lon_bnds.shape[1], 
                                                 lon_bnds.shape[2])
    grid_corner_lat = deepcopy(lat_bnds).reshape(lat_bnds.shape[0]*lat_bnds.shape[1], 
                                                 lat_bnds.shape[2])
    
    # Calculate cell areas
    grid_area = calculate_cell_area(grid_corner_lon, grid_corner_lat)

    return grid_area


def calculate_geometry_area(geometry_list):
    """
    Get coordinate bounds and call function to calculate the area of each cell of a set of geometries.

    Parameters
    ----------
    geometry_list : list
        List with polygon geometries.
    """
    
    # Get coordinate bounds from polygons
    geometry_corner_lon, geometry_corner_lat = get_bnds_from_polygon(geometry_list)

    # Calculate cell areas
    geometry_area = calculate_cell_area(geometry_corner_lon, geometry_corner_lat)
    
    return geometry_area


def get_bnds_from_polygon(geometry_list):

    # Extract the point values that define the perimeter of the polygon
    lon_bnds, lat_bnds = geometry_list[0].exterior.coords.xy
    for poly_ind in range(1, len(geometry_list)):
        poly_lon_bnds, poly_lat_bnds = geometry_list[poly_ind].exterior.coords.xy
        lon_bnds = np.vstack((lon_bnds, np.array(poly_lon_bnds)))
        lat_bnds = np.vstack((lat_bnds, np.array(poly_lat_bnds)))

    return lon_bnds, lat_bnds


def calculate_cell_area(grid_corner_lon, grid_corner_lat):
    """
    Calculate the area of each cell of a grid.

    Parameters
    ----------
    lon_bnds : np.array
        Array with longitude bounds of grid.
    lat_bnds : np.array
        Array with longitude bounds of grid.
    """

    # Calculate area for each grid cell
    n_cells = grid_corner_lon.shape[0]
    area = np.empty(shape=(n_cells,))
    for i in range(0, n_cells):
        area[i] = mod_huiliers_area(grid_corner_lon[i], grid_corner_lat[i]) 

    # Get minor and major axes of Earth
    # TODO: Create file to store dict and call
    earth_radius = {'WGS84': [6356752.3142, 6378137.0]}

    return area*earth_radius['WGS84'][0]*earth_radius['WGS84'][1]


def mod_huiliers_area(cell_corner_lon, cell_corner_lat):
    """
    Calculate the area of each cell according to Huilier's theorem.
    Reference: CDO (https://earth.bsc.es/gitlab/ces/cdo/)

    Parameters
    ----------
    cell_corner_lon : np.array
        Longitude boundaries of each cell
    cell_corner_lat : np.array
        Latitude boundaries of each cell
    """

    sum = 0

    # Get points 0 (bottom left) and 1 (bottom right) in Earth coordinates
    point_0 = lon_lat_to_cartesian(cell_corner_lon[0], cell_corner_lat[0], radius=1)
    point_1 = lon_lat_to_cartesian(cell_corner_lon[1], cell_corner_lat[1], radius=1)
    point_0, point_1 = point_0[0], point_1[0]
    
    # Get number of vertices
    if cell_corner_lat[0] == cell_corner_lat[-1]:
        spatial_nv = len(cell_corner_lon) - 1
    else:
        spatial_nv = len(cell_corner_lon)
    
    for i in range(2, spatial_nv):
        
        # Get point 2 (top right) in Earth coordinates
        point_2 = lon_lat_to_cartesian(cell_corner_lon[i], cell_corner_lat[i], radius=1)
        point_2 = point_2[0]

        # Calculate area of triangle between points 0, 1 and 2
        sum += tri_area(point_0, point_1, point_2)
        
        # Copy to calculate area of next triangle
        if i == (spatial_nv -  1):
            point_1 = deepcopy(point_2)
    
    return sum


def tri_area(point_0, point_1, point_2):
    """
    Calculate area between three points that form a triangle.
    Reference: CDO (https://earth.bsc.es/gitlab/ces/cdo/)

    Parameters
    ----------
    point_0 : np.array
        Position of first point in cartesian coordinates.
    point_1 : np.array
        Position of second point in cartesian coordinates.
    point_2 : np.array
        Position of third point in cartesian coordinates.
    """

    # Get length of side a (between point 0 and 1)
    tmp_vec = cross_product(point_0, point_1)
    sina = norm(tmp_vec)
    a = np.arcsin(sina)

    # Get length of side b (between point 0 and 2)
    tmp_vec = cross_product(point_0, point_2)
    sinb = norm(tmp_vec)
    b = np.arcsin(sinb)

    # Get length of side c (between point 1 and 2)
    tmp_vec = cross_product(point_2, point_1)
    sinc = norm(tmp_vec)
    c = np.arcsin(sinc)

    # Calculate area
    s = 0.5*(a+b+c)
    t = np.tan(s*0.5) * np.tan((s - a)*0.5) * np.tan((s - b)*0.5) * np.tan((s - c)*0.5)
    area = np.fabs(4.0 * np.arctan(np.sqrt(np.fabs(t))))
    
    return area


def cross_product(a, b):
    """
    Calculate cross product between two points.

    Parameters
    ----------
    a : np.array
        Position of point a in cartesian coordinates.
    b : np.array
        Position of point b in cartesian coordinates.
    """

    return [a[1]*b[2] - a[2]*b[1], 
            a[2]*b[0] - a[0]*b[2], 
            a[0]*b[1] - a[1]*b[0]]


def norm(cp):
    """
    Normalize the result of the cross product operation.

    Parameters
    ----------
    cp : np.array
        Cross product between two points.
    """

    return np.sqrt(cp[0]*cp[0] + cp[1]*cp[1] + cp[2]*cp[2])


def lon_lat_to_cartesian(lon, lat, radius=None):
    """
    Calculate lon, lat coordinates of a point on a sphere.

    Parameters
    ----------
    lon : np.array
        Longitude values.
    lat : np.array
        Latitude values.
    radius : float
        Radius of the sphere to get the distances.
    """

    if radius is None:
        # Get major axis of Earth
        # TODO: Create file to store dict and call
        earth_radius = {'WGS84': [6356752.3142, 6378137.0]}
        radius = earth_radius['WGS84'][1]

    lon_r = np.radians(lon)
    lat_r = np.radians(lat)

    x = radius * np.cos(lat_r) * np.cos(lon_r)
    y = radius * np.cos(lat_r) * np.sin(lon_r)
    z = radius * np.sin(lat_r)

    return np.column_stack([x, y, z])
