#!/usr/bin/env python

import sys
import os
from venv import create
from mpi4py import MPI
from netCDF4 import Dataset
from nes.nc_projections.default_nes import Nes
from nes.nc_projections.rotated_nes import RotatedNes
from nes.nc_projections.latlon_nes import LatLonNes
from nes.nc_projections.points_nes import PointsNes
from nes.nc_projections.points_nes_ghost import PointsNesGHOST

def open_netcdf(path, comm=None, xarray=False, info=False, parallel_method=None,
                avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None,
                create_nes=False):
    """
    Open a netCDF file

    Parameters
    ----------
    path : str
        Path to the NetCDF file to read
    comm : MPI.COMM
        MPI communicator to use in that netCDF. Default: MPI.COMM_WORLD
    xarray : bool
        (Not working) Indicates if you want to use xarray. Default: False
    info : bool
        Indicates if you want to print (stdout) the reading/writing steps
    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.
    parallel_method : str
        Indicates the parallelization method that you want. Default: 'Y' (over Y axis)
        accepted values: ['X', 'Y', 'T']

    Returns
    -------
    Nes
        Nes object. Variables read in lazy mode (only metadata)
    """
    if comm is None:
        comm = MPI.COMM_WORLD
    else:
        comm = comm

    if xarray:
        dataset = None
    else:
        if comm.Get_size() == 1:
            dataset = Dataset(path, format="NETCDF4", mode='r', parallel=False)
        else:
            dataset = Dataset(path, format="NETCDF4", mode='r', parallel=True, comm=comm, info=MPI.Info())

    if __is_rotated(dataset):
        # Rotated grids
        nessy = RotatedNes(comm=comm, dataset=dataset, xarray=xarray, 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, create_nes=create_nes)
    elif __is_points(dataset):
        if __is_points_GHOST(dataset):
            # Points - GHOST
            nessy = PointsNesGHOST(comm=comm, dataset=dataset, xarray=xarray, 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, create_nes=create_nes)
        else:
            # Points - non-GHOST
            nessy = PointsNes(comm=comm, dataset=dataset, xarray=xarray, 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, create_nes=create_nes)
    else:
        # Regular grids
        nessy = LatLonNes(comm=comm, dataset=dataset, xarray=xarray, 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, create_nes=create_nes)

    return nessy

def __is_rotated(dataset):
    """
    Check if the netCDF is in rotated pole projection or not.

    Parameters
    ----------
    dataset : Dataset
        netcdf4-python opened dataset object

    Returns
    -------
    value : bool
        Indicated if the netCDF is a rotated one
    """
    if 'rotated_pole' in dataset.variables.keys():
        return True
    else:
        return False

def __is_points(dataset):

    if 'station' in dataset.dimensions:
        return True
    else:
        return False

def __is_points_GHOST(dataset):

    if 'N_flag_codes' in dataset.dimensions and 'N_qa_codes' in dataset.dimensions:
        return True
    else:
        return False
        