#!/usr/bin/env python

import warnings
from netCDF4 import num2date
from mpi4py import MPI
import geopandas as gpd
from .nc_projections import *


def create_nes(comm=None, info=False, projection=None, parallel_method='Y', balanced=False, 
               strlen=75, times=None, avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, 
               **kwargs):
    """
    Create a Nes class from scratch.

    Parameters
    ----------
    comm: MPI.Communicator
        MPI Communicator.
    info: bool
        Indicates if you want to get reading/writing info.
    parallel_method : str
        Indicates the parallelization method that you want. Default: 'Y'.
        accepted values: ['X', 'Y', 'T'].
    balanced : bool
        Indicates if you want a balanced parallelization or not. 
        Balanced dataset cannot be written in chunking mode.
    avoid_first_hours : int
        Number of hours to remove from first time steps.
    avoid_last_hours : int
        Number of hours to remove from last time steps.
    first_level : int
        Index of the first level to use.
    last_level : int, None
        Index of the last level to use. None if it is the last.
    kwargs :
        Projection dependent parameters to create it from scratch.
    """
        
    if comm is None:
        comm = MPI.COMM_WORLD
    else:
        comm = comm

    # Create time array
    if times is None:
        units = 'days since 1996-12-31 00:00:00'
        calendar = 'standard'
        times = num2date([0], units=units, calendar=calendar)
        times = [aux.replace(second=0, microsecond=0) for aux in times]
    else:
        if not isinstance(times, list):
            times = times.tolist()

    # Check if the parameters that are required to create the object have been defined in kwargs
    kwargs_list = []
    for name, value in kwargs.items():
        kwargs_list.append(name)

    if projection is None:
        required_vars = ['lat', 'lon']
    elif projection == 'regular':
        required_vars = ['lat_orig', 'lon_orig', 'inc_lat', 'inc_lon', 'n_lat', 'n_lon']
    elif projection == 'rotated':
        required_vars = ['centre_lat', 'centre_lon', 'west_boundary', 'south_boundary', 'inc_rlat', 'inc_rlon']
    elif projection == 'lcc':
        required_vars = ['lat_1', 'lat_2', 'lon_0', 'lat_0', 'nx', 'ny', 'inc_x', 'inc_y', 'x_0', 'y_0']
    elif projection == 'mercator':
        required_vars = ['lat_ts', 'lon_0', 'nx', 'ny', 'inc_x', 'inc_y', 'x_0', 'y_0']
    else:
        raise ValueError("Unknown projection: {0}".format(projection))

    for var in required_vars:
        if var not in kwargs_list:
            msg = 'Variable {0} has not been defined. '.format(var)
            msg += 'For a {} projection, it is necessary to define {}'.format(projection, required_vars)
            raise ValueError(msg)

    if projection is None:
        if parallel_method == 'Y':
            warnings.warn("Parallel method cannot be 'Y' to create points NES. Setting it to 'X'")
            parallel_method = 'X'
        elif parallel_method == 'T':
            raise NotImplementedError("Parallel method T not implemented yet")
        nessy = PointsNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method,
                          avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours,
                          first_level=first_level, last_level=last_level, balanced=balanced,
                          create_nes=True, strlen=strlen, times=times, **kwargs)
    elif projection == 'regular':
        nessy = LatLonNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method,
                          avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours,
                          first_level=first_level, last_level=last_level, balanced=balanced,
                          create_nes=True, times=times, **kwargs)
    elif projection == 'rotated':
        nessy = RotatedNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method,
                           avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours,
                           first_level=first_level, last_level=last_level, balanced=balanced,
                           create_nes=True, times=times, **kwargs)
    elif projection == 'lcc':
        nessy = LCCNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method,
                       avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours,
                       first_level=first_level, last_level=last_level, balanced=balanced,
                       create_nes=True, times=times, **kwargs)
    elif projection == 'mercator':
        nessy = MercatorNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method,
                            avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours,
                            first_level=first_level, last_level=last_level, balanced=balanced,
                            create_nes=True, times=times, **kwargs)
    else:
        raise NotImplementedError(projection)
    
    return nessy
    

def from_shapefile(path, method=None, **kwargs):
    """
    Create NES from shapefile data.
    
    1. Create NES grid.
    2. Create shapefile for grid.
    3. Read shapefile from mask.
    4. Spatial join to add shapefile variables to NES variables.

    Parameters
    ----------
    path : str
        Path to shapefile.
    kwargs :
        Projection and projection dependent parameters to create it from scratch.

    method : str
        Overlay method. Accepted values: ['nearest', 'intersection', None].
    """
    
    # Create NES
    nessy = create_nes(comm=None, info=False, **kwargs)

    # Create shapefile for grid
    nessy.create_shapefile()

    # Read shapefile
    mask = gpd.read_file(path)

    # Make spatial join
    nessy.spatial_join(mask, method=method)

    return nessy
