#!/usr/bin/env python
from nes import create_nes, open_netcdf
from hermesv3_gr.config import log_message, get_time_stamp, record_time
import os
import pytz
from datetime import timedelta
from datetime import timezone as dtz
import pandas as pd

BALANCED = True


def select_grid(comm, options, times=None):
    """
    Select & create the destination grid.

    1. If grid auxiliary file exists will be read, otherwise will be created


    Parameters
    ----------
    comm : MPI.COMMUNICATOR
        MPI communicator.
    options : NameSpace
        Configuration options.
    times : List[datetime]
        List of date times

    Returns
    -------
    Nes
        A NES grid object with the shapefile & cell_area calculated.
    """
    st_time = get_time_stamp()
    grid_path = os.path.join(options.auxiliary_files_path, 'grid.nc')

    if os.path.exists(grid_path):
        log_message('Reading existing grid from {0}'.format(grid_path), level=1)
        grid = open_netcdf(grid_path, comm=comm, balanced=BALANCED)
        grid.load()
        write_grid = False
    else:
        log_message('Creating grid', level=1)
        write_grid = True

        # Creating a different object depending on the grid type
        if options.domain_type == 'global':
            log_message("Global grid with:" +
                               "\n\t\tinc_lat: {0}\n\t\tinc_lon: {1}".format(options.inc_lat, options.inc_lon),
                               level=1)
            grid = create_nes(comm=comm, info=False, times=times, projection='global',
                              inc_lat=options.inc_lat, inc_lon=options.inc_lon, balanced=BALANCED)

        elif options.domain_type == 'regular':
            log_message("Creating Regular Lat-Lon grid with:" +
                               "\n\t\tlat_orig: {0}\n\t\tlon_orig: {1}".format(options.lat_orig, options.lon_orig) +
                               "\n\t\tinc_lat: {0}\n\t\tinc_lon: {1}".format(options.inc_lat, options.inc_lon) +
                               "\n\t\tn_lat: {0}\n\t\tn_lon: {1}".format(options.n_lat, options.n_lon),
                               level=1)
            grid = create_nes(comm=comm, info=False, times=times, projection='regular',
                              lat_orig=options.lat_orig, lon_orig=options.lon_orig, inc_lat=options.inc_lat,
                              inc_lon=options.inc_lon, n_lat=options.n_lat, n_lon=options.n_lon, balanced=BALANCED)
            log_message('Regular grid created', level=2)

        elif options.domain_type == 'rotated':
            log_message("Creating Rotated grid selected with:" +
                               "\n\t\tcentre_lat: {0}\n\t\tcentre_lon: {1}".format(options.centre_lat,
                                                                                   options.centre_lon) +
                               "\n\t\twest_boundary: {0}\n\t\tsouth_boundary: {1}".format(options.west_boundary,
                                                                                          options.south_boundary) +
                               "\n\t\tinc_rlat: {0}\n\t\tinc_rlon: {1}".format(options.inc_rlat, options.inc_rlon),
                               level=1)
            grid = create_nes(comm=comm, info=False, times=times, projection='rotated',
                              centre_lat=options.centre_lat, centre_lon=options.centre_lon,
                              west_boundary=options.west_boundary, south_boundary=options.south_boundary,
                              inc_rlat=options.inc_rlat, inc_rlon=options.inc_rlon, balanced=BALANCED)

        elif options.domain_type == 'rotated_nested':
            log_message("Creating Rotated-Nested grid with:" +
                               "\n\t\tparent_grid_path: {0}".format(options.parent_grid_path) +
                               "\n\t\tparent_ratio: {0}".format(options.parent_ratio) +
                               "\n\t\ti_parent_start: {0}".format(options.i_parent_start) +
                               "\n\t\tj_parent_start: {0}".format(options.j_parent_start) +
                               "\n\t\tn_rlat: {0}\n\t\tn_rlat: {1}".format(options.n_rlat, options.n_lon),
                               level=1)
            grid = create_nes(comm=comm, info=False, times=times, projection='rotated-nested',
                              parent_grid_path=options.parent_grid_path, parent_ratio=options.parent_ratio,
                              i_parent_start=options.i_parent_start, j_parent_start=options.j_parent_start,
                              n_rlat=options.n_rlat, n_rlon=options.n_rlon, balanced=BALANCED)

        elif options.domain_type == 'lcc':
            log_message("Creating Lambert Conformal Conic grid with:" +
                               "\n\t\tlat_1: {0}\n\t\tlat_2: {1}".format(options.lat_1, options.lat_2) +
                               "\n\t\tlon_0: {0}\n\t\tlat_0: {1}".format(options.lon_0, options.lat_0) +
                               "\n\t\tnx: {0}\n\t\tny: {1}".format(options.nx, options.ny) +
                               "\n\t\tinc_x: {0}\n\t\tinc_y: {1}".format(options.inc_x, options.inc_y) +
                               "\n\t\tx_0: {0}\n\t\ty_0: {1}".format(options.x_0, options.y_0),
                               level=1)
            grid = create_nes(comm=comm, info=False, times=times, projection='lcc',
                              lat_1=options.lat_1, lat_2=options.lat_2, lon_0=options.lon_0, lat_0=options.lat_0,
                              nx=options.nx, ny=options.ny, inc_x=options.inc_x, inc_y=options.inc_y,
                              x_0=options.x_0, y_0=options.y_0, balanced=BALANCED)

        elif options.domain_type == 'mercator':
            log_message("Creating Mercator grid with:" +
                               "\n\t\tlat_ts: {0}\n\t\tlon_0: {1}".format(options.lat_ts, options.lon_0) +
                               "\n\t\tnx: {0}\n\t\tny: {1}".format(options.nx, options.ny) +
                               "\n\t\tinc_x: {0}\n\t\tinc_y: {1}".format(options.inc_x, options.inc_y) +
                               "\n\t\tx_0: {0}\n\t\ty_0: {1}".format(options.x_0, options.y_0),
                               level=1)
            grid = create_nes(comm=comm, info=False, times=times, projection='mercator',
                              lat_ts=options.lat_ts, lon_0=options.lon_0, nx=options.nx, ny=options.ny,
                              inc_x=options.inc_x, inc_y=options.inc_y, x_0=options.x_0, y_0=options.y_0,
                              balanced=BALANCED)
        else:
            log_message('ERROR: Check the .err file to get more info.')
            raise NotImplementedError("The grid type {0} is not implemented.".format(options.domain_type)
                                      + " Use 'global', 'regular, 'rotated', 'rotated_nested', 'lcc' or 'mercator.")
    # Cell area
    if 'cell_area' not in grid.cell_measures.keys():
        write_grid = True
        grid.calculate_grid_area()
        log_message('Grid: cell area created', level=3)

    # Shapefile
    grid.create_shapefile()
    log_message('Grid: shapefile created', level=3)

    # Timezones
    if options.timezones_shapefile is None:
        raise FileNotFoundError("No timezones shapefile was defined. Fill the '--timezones_shapefile' option.")

    if 'tzid' not in grid.variables.keys():
        write_grid = True
        log_message('Grid: Timezones calculation starting', level=3)
        grid.spatial_join(options.timezones_shapefile, method='nearest')
        grid.variables['tzid'] = {
            'data': grid.shapefile['tzid'].values.reshape((grid.lat['data'].shape[0], grid.lon['data'].shape[-1])),
            'dtype': str}
        grid.set_strlen(32)
        log_message('Grid: Timezones created', level=3)
    else:
        grid.shapefile['tzid'] = grid.variables['tzid']['data'].flatten()
        log_message('Grid: Timezones read from grid.nc', level=3)

    # Write
    if write_grid:
        grid.to_netcdf(grid_path, serial=True)
        log_message('Grid: file created: {0}'.format(grid_path), level=2)

    log_message('Destination grid done!', level=2)

    record_time('Grid', 'select_grid', get_time_stamp() - st_time)
    return grid


def add_local_dates(grid, date_time, pandas=False):
    def parse_tz(timezone):
        """
        Parse the timezone (string format).

        It is needed because some libraries have more timezones than others, and it
        tries to simplify setting the strange ones into the nearest common one.
        Examples:
            'America/Punta_Arenas': 'America/Santiago',
            'Europe/Astrakhan': 'Europe/Moscow',
            'Asia/Atyrau': 'Asia/Aqtau',
            'Asia/Barnaul': 'Asia/Almaty',
            'Europe/Saratov': 'Europe/Moscow',
            'Europe/Ulyanovsk': 'Europe/Moscow',
            'Europe/Kirov': 'Europe/Moscow',
            'Asia/Tomsk': 'Asia/Novokuznetsk',
            'America/Fort_Nelson': 'America/Vancouver'

        Parameters
        ----------
        timezone : str
            Not parsed timezone.

        Returns
        -------
        str
            Parsed timezone
        """

        tz_dict = {
            'America/Punta_Arenas': 'America/Santiago',
            'Europe/Astrakhan': 'Europe/Moscow',
            'Asia/Atyrau': 'Asia/Aqtau',
            'Asia/Barnaul': 'Asia/Almaty',
            'Europe/Saratov': 'Europe/Moscow',
            'Europe/Ulyanovsk': 'Europe/Moscow',
            'Europe/Kirov': 'Europe/Moscow',
            'Asia/Tomsk': 'Asia/Novokuznetsk',
            'America/Fort_Nelson': 'America/Vancouver',
            'Asia/Famagusta': 'Asia/Nicosia',
            'America/Nuuk': dtz(timedelta(hours=-3)),
        }

        if timezone in iter(tz_dict.keys()):
            timezone = tz_dict[timezone]

        return timezone
    grid.shapefile['local'] = pd.to_datetime(date_time, utc=True)
    try:
        # print(f"{grid.rank}, real local?: {grid.shapefile.groupby('tzid')['local'].apply(
        #     lambda x: x.dt.tz_convert(x.name).dt.tz_localize(None)).reset_index(level='tzid', drop=True)}")
        grid.shapefile['local'] = grid.shapefile.groupby('tzid')['local'].apply(
            lambda x: x.dt.tz_convert(x.name).dt.tz_localize(None)).reset_index(level='tzid', drop=True)
    # except TypeError as e:
    #     print(f"{grid.rank}, date_time: {date_time}")
    #     print(f"{grid.rank}, grid.shapefile['local']: {grid.shapefile['local']}")
    #     print(f"{grid.rank}, grid.shapefile['tzid']: {grid.shapefile['tzid']}")
    #     raise e
    except pytz.exceptions.UnknownTimeZoneError:
        grid.shapefile['local'] = grid.shapefile.groupby('tzid')['local'].apply(
            lambda x: x.dt.tz_convert(parse_tz(x.name)).dt.tz_localize(None))
    if pandas:
        time_stamp = grid.shapefile['local'].copy()
    else:
        time_stamp = grid.shapefile['local'].to_numpy().reshape((grid.lat['data'].shape[0], grid.lon['data'].shape[-1]))

    del grid.shapefile['local']

    return time_stamp
