#!/usr/bin/env python

import os
import sys
from mpi4py import MPI
from netCDF4 import Dataset
import warnings
import numpy as np
from .nc_projections import *

DIM_VAR_NAMES = ['lat', 'latitude', 'lat_bnds', 'lon', 'longitude', 'lon_bnds', 'time', 'time_bnds', 'lev', 'level',
                 'cell_area', 'crs', 'rotated_pole', 'x', 'y', 'rlat', 'rlon', 'Lambert_conformal', 'mercator']


def open_netcdf(path, comm=None, info=False, parallel_method='Y', avoid_first_hours=0, avoid_last_hours=0,
                first_level=0, last_level=None, balanced=False):
    """
    Open a NetCDF file and return a Nes object with lazily loaded variables.

    Parameters
    ----------
    path : str
        Path to the NetCDF file to read.
    comm : MPI.COMM, optional
        MPI communicator to use for parallel I/O. Default: MPI.COMM_WORLD.
    info : bool, optional
        Indicates whether to print reading/writing steps to stdout.
    avoid_first_hours : int, optional
        Number of hours to remove from the beginning of time steps.
    avoid_last_hours : int, optional
        Number of hours to remove from the end of time steps.
    parallel_method : str, optional
        Parallelization method. Default: 'Y'.
        Accepted values: ['X', 'Y', 'T']
    balanced : bool, optional
        Indicates if you want balanced parallelization. Note: Balanced datasets cannot be written in chunking mode.
    first_level : int, optional
        Index of the first level to use.
    last_level : int or None, optional
        Index of the last level to use. None if it is the last.

    Returns
    -------
    Nes
        Nes object containing lazily loaded variables with metadata.

    Raises
    ------
    FileNotFoundError
        If the specified NetCDF file at `path` does not exist.

    Notes
    -----
    This function supports parallel I/O using MPI. The `comm` parameter allows you to specify an MPI communicator for coordinating processes.

    Examples
    --------
    Open a NetCDF file using the default communicator:

    ```python
    nes_data = open_netcdf('/path/to/data.nc', info=True, parallel_method='Y')
    ```

    Open a NetCDF file with a custom MPI communicator:

    ```python
    from mpi4py import MPI

    comm = MPI.COMM_WORLD
    nes_data = open_netcdf('/path/to/data.nc', comm=comm, info=True, parallel_method='Y')
    ```

    Notes
    -----
    - If `comm` is not provided, the default MPI communicator `MPI.COMM_WORLD` is used.
    - The `info` parameter controls whether to print reading/writing steps to the standard output.

    """
    # Set default communicator if not provided
    comm = comm or MPI.COMM_WORLD
    # Validate file existence
    if not os.path.exists(path):
        raise FileNotFoundError(f"The specified NetCDF file at '{path}' does not exist.")

    # Open NetCDF dataset
    dataset = Dataset(path, format="NETCDF4", mode='r', parallel=False)

    # Determine grid type and create Nes object
    if __is_rotated(dataset):
        # Rotated grids
        nessy = RotatedNes(comm=comm, dataset=dataset, 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=False, balanced=balanced)
    elif __is_points(dataset):
        if parallel_method == 'Y':
            warnings.warn("Parallel method cannot be 'Y' to create points NES. Setting it to 'X'")
            sys.stderr.flush()
            parallel_method = 'X'
        if __is_points_ghost(dataset):
            # Points - GHOST
            nessy = PointsNesGHOST(comm=comm, dataset=dataset, 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=False, balanced=balanced)
        elif __is_points_providentia(dataset):
            # Points - Providentia
            nessy = PointsNesProvidentia(comm=comm, dataset=dataset, 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=False,
                                         balanced=balanced)
        else:
            # Points - non-GHOST
            nessy = PointsNes(comm=comm, dataset=dataset, 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=False, balanced=balanced)
    elif __is_lcc(dataset):
        # Lambert conformal conic grids
        nessy = LCCNes(comm=comm, dataset=dataset, 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=False, balanced=balanced)
    elif __is_mercator(dataset):
        # Mercator grids
        nessy = MercatorNes(comm=comm, dataset=dataset, 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=False, balanced=balanced)
    else:
        # Regular grids
        nessy = LatLonNes(comm=comm, dataset=dataset, 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=False, balanced=balanced)

    return nessy


def __is_rotated(dataset):
    """
    Check if the NetCDF dataset uses a rotated pole projection.

    Parameters
    ----------
    dataset : netCDF4.Dataset
        An open NetCDF dataset.

    Returns
    -------
    bool
        True if the dataset uses a rotated pole projection, False otherwise.

    Notes
    -----
    This function checks for the presence of either a 'rotated_pole' variable or
    dimensions 'rlat' and 'rlon' in the NetCDF dataset to determine if it is in a
    rotated pole projection.
    """
    rotated_pole_variable = 'rotated_pole' in dataset.variables
    rlat_rlon_dimensions = 'rlat' in dataset.dimensions and 'rlon' in dataset.dimensions

    return rotated_pole_variable or rlat_rlon_dimensions


def __is_points(dataset):
    """
    Check if the NetCDF dataset represents a non-GHOST points dataset.

    Parameters
    ----------
    dataset : netCDF4.Dataset
        An open NetCDF dataset.

    Returns
    -------
    bool
        True if the dataset represents a non-GHOST points dataset, False otherwise.

    Notes
    -----
    This function checks for the presence of a 'station' dimension in the NetCDF dataset
    to determine if it represents a points dataset in non-GHOST format.
    """
    return 'station' in dataset.dimensions


def __is_points_ghost(dataset):
    """
    Check if the NetCDF dataset represents a GHOST points dataset.

    Parameters
    ----------
    dataset : netCDF4.Dataset
        An open NetCDF dataset.

    Returns
    -------
    bool
        True if the dataset represents a GHOST points dataset, False otherwise.

    Notes
    -----
    This function checks for the presence of 'N_flag_codes' and 'N_qa_codes' dimensions
    in the NetCDF dataset to determine if it represents a points dataset in GHOST format.

    """
    return 'N_flag_codes' in dataset.dimensions and 'N_qa_codes' in dataset.dimensions


def __is_points_providentia(dataset):
    """
    Check if the NetCDF dataset represents a Providentia points dataset.

    Parameters
    ----------
    dataset : netCDF4.Dataset
        An open NetCDF dataset.

    Returns
    -------
    bool
        True if the dataset represents a Providentia points dataset, False otherwise.

    Notes
    -----
    This function checks for the presence of 'grid_edge', 'model_latitude', and 'model_longitude'
    dimensions in the NetCDF dataset to determine if it represents a points dataset in Providentia format.

    """
    return ('grid_edge' in dataset.dimensions and 'model_latitude' in dataset.dimensions and
            'model_longitude' in dataset.dimensions)


def __is_lcc(dataset):
    """
    Check if the NetCDF dataset is in Lambert Conformal Conic (LCC) projection.

    Parameters
    ----------
    dataset : netCDF4.Dataset
        An open NetCDF dataset.

    Returns
    -------
    bool
        True if the dataset is in Lambert Conformal Conic (LCC) projection, False otherwise.

    Notes
    -----
    This function checks for the presence of 'Lambert_Conformal' or 'Lambert_conformal'
    variables in the NetCDF dataset to determine if it is in Lambert Conformal Conic projection.
    """
    return 'Lambert_Conformal' in dataset.variables or 'Lambert_conformal' in dataset.variables


def __is_mercator(dataset):
    """
    Check if the NetCDF dataset is in Mercator projection.

    Parameters
    ----------
    dataset : netCDF4.Dataset
        An open NetCDF dataset.

    Returns
    -------
    bool
        True if the dataset is in Mercator projection, False otherwise.

    Notes
    -----
    This function checks for the presence of 'mercator' variable in the NetCDF dataset
    to determine if it is in Mercator projection.

    """
    return 'mercator' in dataset.variables


def concatenate_netcdfs(nessy_list, comm=None, info=False, parallel_method='Y', avoid_first_hours=0, avoid_last_hours=0,
                        first_level=0, last_level=None, balanced=False):
    """
    Concatenate variables from different sources.

    Parameters
    ----------
    nessy_list : list of Nes or str
        List of Nes objects or paths to NetCDF files for concatenation.
    comm : MPI.Communicator, optional
        MPI Communicator. Default: None.
    info : bool, optional
        Indicates whether to print reading/writing steps to stdout. Default: False.
    parallel_method : str, optional
        Parallelization method. Default: 'Y'. Accepted values: ['X', 'Y', 'T']
    avoid_first_hours : int, optional
        Number of hours to remove from the beginning of time steps. Default: 0.
    avoid_last_hours : int, optional
        Number of hours to remove from the end of time steps. Default: 0.
    first_level : int, optional
        Index of the first level to use. Default: 0.
    last_level : int or None, optional
        Index of the last level to use. None if it is the last. Default: None.
    balanced : bool, optional
        Indicates if balanced parallelization is desired. Note: Balanced datasets cannot be written in chunking mode.
        Default: False.

    Returns
    -------
    Nes
        Nes object with all the concatenated variables.

    Raises
    ------
    TypeError
        If the input is not a list.
    FileNotFoundError
        If a provided path to a NetCDF file does not exist.
    ValueError
        If the input list is empty.
    """
    if not isinstance(nessy_list, list):
        raise TypeError("You must pass a list of NES objects or paths.")

    if not nessy_list:
        raise ValueError("Input list is empty.")

    if isinstance(nessy_list[0], str):
        if not os.path.exists(nessy_list[0]):
            raise FileNotFoundError(nessy_list[0])
        nessy_first = open_netcdf(nessy_list[0],
                                  comm=comm,
                                  parallel_method=parallel_method,
                                  info=info,
                                  avoid_first_hours=avoid_first_hours,
                                  avoid_last_hours=avoid_last_hours,
                                  first_level=first_level,
                                  last_level=last_level,
                                  balanced=balanced
                                  )
        nessy_first.load()
    else:
        nessy_first = nessy_list[0]
    for aux_nessy in nessy_list[1:]:
        if isinstance(aux_nessy, str):
            if not os.path.exists(aux_nessy):
                raise FileNotFoundError(aux_nessy)
            nc_add = Dataset(filename=aux_nessy, mode='r')
            for var_name, var_info in nc_add.variables.items():
                if var_name not in DIM_VAR_NAMES:
                    nessy_first.variables[var_name] = {}
                    var_dims = var_info.dimensions
                    # Read data in 4 dimensions
                    if len(var_dims) < 2:
                        data = var_info[:]
                    elif len(var_dims) == 2:
                        data = var_info[nessy_first.read_axis_limits['y_min']:nessy_first.read_axis_limits['y_max'],
                                        nessy_first.read_axis_limits['x_min']:nessy_first.read_axis_limits['x_max']]
                        data = data.reshape(1, 1, data.shape[-2], data.shape[-1])
                    elif len(var_dims) == 3:
                        if 'strlen' in var_dims:
                            data = var_info[nessy_first.read_axis_limits['y_min']:nessy_first.read_axis_limits['y_max'],
                                            nessy_first.read_axis_limits['x_min']:nessy_first.read_axis_limits['x_max'],
                                            :]
                            data_aux = np.empty(shape=(data.shape[0], data.shape[1]), dtype=np.object)
                            for lat_n in range(data.shape[0]):
                                for lon_n in range(data.shape[1]):
                                    data_aux[lat_n, lon_n] = ''.join(
                                        data[lat_n, lon_n].tostring().decode('ascii').replace('\x00', ''))
                            data = data_aux.reshape((1, 1, data_aux.shape[-2], data_aux.shape[-1]))
                        else:
                            data = var_info[nessy_first.read_axis_limits['t_min']:nessy_first.read_axis_limits['t_max'],
                                            nessy_first.read_axis_limits['y_min']:nessy_first.read_axis_limits['y_max'],
                                            nessy_first.read_axis_limits['x_min']:nessy_first.read_axis_limits['x_max']]
                            data = data.reshape(data.shape[-3], 1, data.shape[-2], data.shape[-1])
                    elif len(var_dims) == 4:
                        data = var_info[nessy_first.read_axis_limits['t_min']:nessy_first.read_axis_limits['t_max'],
                                        nessy_first.read_axis_limits['z_min']:nessy_first.read_axis_limits['z_max'],
                                        nessy_first.read_axis_limits['y_min']:nessy_first.read_axis_limits['y_max'],
                                        nessy_first.read_axis_limits['x_min']:nessy_first.read_axis_limits['x_max']]
                    else:
                        raise TypeError("{} data shape is nto accepted".format(var_dims))

                    nessy_first.variables[var_name]['data'] = data
                    # Avoid some attributes
                    for attrname in var_info.ncattrs():
                        if attrname not in ['missing_value', '_FillValue']:
                            value = getattr(var_info, attrname)
                            if value in ['unitless', '-']:
                                value = ''
                            nessy_first.variables[var_name][attrname] = value
            nc_add.close()

        else:
            nessy_first.concatenate(aux_nessy)
        
    return nessy_first


def open_raster(path, comm=None, info=False, parallel_method='Y'):
    """
    Open a Raster (TIFF) file and return a Raster NES object.

    This function reads a Raster (TIFF) file, initializes a Raster NES object, and returns it.

    Parameters
    ----------
    path : str
        The path to the Raster (TIFF) file to be read.

    comm : MPI.COMM, optional
        MPI communicator to use in the Raster NES object. Default: MPI.COMM_WORLD.

    info : bool, optional
        Indicates whether to print the reading/writing steps to stdout.

    parallel_method : str, optional
        Specifies the parallelization method to be used in the Raster NES object. Default: 'Y'.
        Accepted values: ['Y', 'X'].

    Returns
    -------
    RasterNes
        The initialized Raster NES object.

    Raises
    ------
    FileNotFoundError
        If the specified Raster (TIFF) file does not exist.

    RuntimeError
        If an unsupported parallelization method ('T') is specified.

    Notes
    -----
    - If the MPI communicator (`comm`) is not provided, it defaults to MPI.COMM_WORLD.
    - The Raster NES object is created using the RasterNes class with the specified parameters.

    Example
    -------
    >>> raster_path = 'path/to/your/raster.tif'
    >>> raster_nes_obj = open_raster(raster_path, comm=comm, info=True, parallel_method='Y')
    """
    raise NotImplementedError('Raster NES object')
    if comm is None:
        comm = MPI.COMM_WORLD

    if not os.path.exists(path):
        raise FileNotFoundError(f"The specified Raster (TIFF) file does not exist: {path}")

    if parallel_method == 'T':
        raise RuntimeError("Parallel method 'T' is not accepted for opening raster files.")

    nessy = RasterNes(path=path, comm=comm, info=info, parallel_method=parallel_method)

    return nessy
