import numpy as np
import pandas as pd
import regionmask
import copy
from typing import Optional, Hashable, Type, List, cast
from canopy import RedSpec, Raster, Field
from canopy.core.grid import get_grid_type
from canopy.core.grid.grid_abc import Grid
from pandas.api.types import is_string_dtype
from canopy.util.checks import check_spatial_coords_match

# Make raster
# -----------
def gcidx(x: float, xmin: float, dx: float) -> int:
    return int((x-(xmin-0.5*dx))/dx)
get_coord_index = np.vectorize(gcidx)

def make_raster(field: Field, layer: str) -> Raster:
    """Produce a Raster object from a Field

    Parameters
    ----------
    field : Field
        The field object from which to create the Raster
    layer : str
        The field layer to rasterize
    """

    data = field.data
    grid = field.grid

    # This is to avoid mypy errors.
    # TODO: find a better solution. Maybe Rasterizable Grid?
    grid_type = get_grid_type(grid) 
    if grid_type != 'lonlat':
        raise ValueError("make_raster currently supports only 'lonlat' grid type.")
    from canopy.core.grid import GridLonLat
    grid  = cast(GridLonLat, grid)

    ilon = get_coord_index(data.index.get_level_values('lon').to_numpy(), grid.lon_min, grid.dlon)
    ilat = get_coord_index(data.index.get_level_values('lat').to_numpy(), grid.lat_min, grid.dlat)

    # Create Raster
    xx, yy = np.meshgrid(grid.lons, grid.lats)

    vmap = np.empty([grid.lats.size, grid.lons.size], dtype=float)
    vmap[:] = np.nan
    dtype = data.dtypes[layer]
    if is_string_dtype(dtype):
        labels = dict(enumerate(data[layer].unique()))
        labels_r = {v: k for (k, v) in labels.items()}
        values = np.vectorize(labels_r.get)(data[layer].values)
        vmap[ilat,ilon] = values
    else:
        labels = None
        vmap[ilat,ilon] = data[layer].values
    
    return Raster(xx, yy, vmap, labels)


# Make lines
# ----------
def make_lines(field: Field, axis: str = 'time', flatten_columns: bool = False, layers: Optional[List[str]] = None) -> pd.DataFrame:
    """Create a pandas DataFrame with the field's data unstacked along the specified axis

    This DataFrame is much easier to use to plot lines (e.g. with df.plot())

    Parameters
    ----------
    field : Field
        The field from which to create the DataFrame
    axis : str
        The axis along which to unstack the data
    flatten_columns : bool, optional
        If True, flatten MultiIndex columns to simple strings. Default is False.
    layers : Optional[List[str]], optional
        Layer names to include when flatten_columns is True. Defaults to field.layers.
    """

    if axis in field.grid.axis_names and field.grid.is_reduced(axis):
        raise ValueError(f"Axis '{axis}' is reduced.")

    if axis == 'label':
        raise ValueError(f"Unstacking along the 'label' axis is not allowed.")

    if axis not in field.data.index.names:
        raise ValueError(f"Axis '{axis}' not found in field's data.")

    df = field._data.copy()
    levels = cast(Hashable, [x for x in df.index.names if not x == axis])
    df = cast(pd.DataFrame, df.unstack(level=levels))

    if flatten_columns:
        layers = layers or field.layers
        df = _flatten_columns(df, layers, field.grid)

    return df


def _flatten_columns(df: pd.DataFrame, layers: List[str], grid: Grid) -> pd.DataFrame:
    """Flatten MultiIndex columns to simple strings."""
    if not isinstance(df.columns, pd.MultiIndex):
        cols = [c for c in df.columns if c in layers]
        return df[cols]

    result_parts = []
    for layer in layers:
        layer_df = df[layer]
        names = layer_df.columns.names

        flat_cols = []
        for row in layer_df.columns:
            if "label" in names:
                flat_cols.append(str(row[names.index("label")]))
            else:
                coord_axes = [ax for ax in grid.axis_names if ax in names]
                values = [row[names.index(ax)] for ax in coord_axes]
                flat_cols.append(
                    f"({', '.join(f'{v:.2f}' for v in values)})" if len(values) > 1 else f"{values[0]:.2f}"
                )

        layer_df.columns = flat_cols
        result_parts.append(layer_df)

    if len(result_parts) == 1:
        return result_parts[0]

    combined = pd.concat(result_parts, axis=1, keys=layers[: len(result_parts)])
    comb_cols = [f"{layer} - {col}" for (layer, col) in combined.columns]
    combined.columns = comb_cols

    return combined


# Filter data by region
# ---------------------

def filter_region(field: Field, region: str, region_type: str = "country") -> Field:
    """Filter a Field to retain only rows whose lon/lat fall inside a named region.
    
    Parameters
    ----------
    field : Field
        The field to filter
    region : str
        Region identifier understood by the chosen region set.
        For example: Names for countries, or region names for Giorgi/SREX/AR6.
    region_type: str
        Which predefined region set to use:
        - "country": natural_earth_v5_0_0.countries_10
        - "giorgi": Giorgi regions
        - "SREX"/"srex": IPCC SREX regions
        - "AR6"/"ar6": IPCC AR6 regions

    Returns
    -------
    A new Field instance containing only rows whose coordinates are inside the region.
    """
    
    # Select which region set to use based on region_type
    region_type_norm = region_type.lower()
    if region_type_norm == "country":
        regionmask_type = regionmask.defined_regions.natural_earth_v5_0_0.countries_10
    elif region_type_norm == "giorgi":
        regionmask_type = regionmask.defined_regions.giorgi
    elif region_type_norm == "srex":
        regionmask_type = regionmask.defined_regions.srex
    elif region_type_norm == "ar6":
        regionmask_type = regionmask.defined_regions.ar6.all
    else:
        raise ValueError("Unsupported region_type; expected one of {'country','giorgi','SREX','AR6'}")

    # Select region
    try:
        region_key = regionmask_type.map_keys(region)
    except KeyError:
        raise ValueError(f"Unknown region '{region}' for region_type '{region_type}'.")
    region_idx = regionmask_type[[region_key]]

    # Compute region mask on the 2D lon-lat grid
    region_mask = region_idx.mask(field.grid.lon, field.grid.lat)
    region_mask = region_mask.assign_coords(lat=field.grid.lat, lon=field.grid.lon)

    # Retrieve field data and coordinates
    data = field.data
    lons = data.index.get_level_values("lon")
    lats = data.index.get_level_values("lat")

    # Get the region_mask coordinate arrays
    mask_lats = region_mask.lat.values
    mask_lons = region_mask.lon.values

    # Find nearest indices (to avoid precision errors)
    lat_indices = np.searchsorted(mask_lats, lats.to_numpy())
    lon_indices = np.searchsorted(mask_lons, lons.to_numpy())

    # Clip to valid range
    lat_indices = np.clip(lat_indices, 0, len(mask_lats) - 1)
    lon_indices = np.clip(lon_indices, 0, len(mask_lons) - 1)

    # Index directly into the region_mask array
    is_in_region = region_mask.values[lat_indices, lon_indices]

    # Add column mask to data
    data["mask"] = np.isfinite(is_in_region) # convert nan to False and 121 to True

    # Apply the boolean mask to keep only rows inside the region
    filtered_field = field.filter("mask == True")

    # Add log message
    filtered_field.log(f"Filter data with region: '{region}'")

    return filtered_field
