import xarray as xr


def _get_feature_importance_array(model_state: xr.Dataset, importance_type: str = "feature_weights") -> xr.DataArray:
    """Get the appropriate feature importance array from model state.

    Args:
        model_state: The xarray Dataset containing the model state.
        importance_type: Type of importance to extract. For ESPA: 'feature_weights'.
            For XGBoost: 'feature_importance_weight', 'feature_importance_gain', etc.
            For RF: 'feature_importance'.

    Returns:
        xr.DataArray with feature importance values.

    """
    if importance_type in model_state:
        return model_state[importance_type]
    # Fallback for compatibility
    if "feature_weights" in model_state:
        return model_state["feature_weights"]
    if "feature_importance" in model_state:
        return model_state["feature_importance"]
    raise ValueError(f"No feature importance array found. Available: {list(model_state.data_vars)}")


def extract_embedding_features(
    model_state: xr.Dataset, importance_type: str = "feature_weights"
) -> xr.DataArray | None:
    """Extract embedding features from the model state.

    Feature naming pattern: `embeddings_{agg}_{band}_{year}`
    Example: `embeddings_mean_B02_2020`

    Args:
        model_state: The xarray Dataset containing the model state.
        importance_type: Type of feature importance to extract.

    Returns:
        xr.DataArray: The extracted embedding features. This DataArray has dimensions
            ('agg', 'band', 'year') corresponding to the different components of the embedding features.
            Returns None if no embedding features are found.

    """

    def _is_embedding_feature(feature: str) -> bool:
        return feature.startswith("embeddings_")

    embedding_features = [f for f in model_state.feature.to_numpy() if _is_embedding_feature(f)]
    if len(embedding_features) == 0:
        return None

    # Get the appropriate importance array
    importance_array = _get_feature_importance_array(model_state, importance_type)

    # Split the single feature dimension of embedding features into separate dimensions (agg, band, year)
    embedding_feature_array = importance_array.sel(feature=embedding_features)
    embedding_feature_array = embedding_feature_array.assign_coords(
        agg=("feature", [f.split("_")[1] for f in embedding_features]),
        band=("feature", [f.split("_")[2] for f in embedding_features]),
        year=("feature", [f.split("_")[3] for f in embedding_features]),
    )
    embedding_feature_array = embedding_feature_array.set_index(feature=["agg", "band", "year"]).unstack("feature")
    return embedding_feature_array


def extract_era5_features(
    model_state: xr.Dataset,
    importance_type: str = "feature_weights",
    temporal_group: str | None = None,
) -> xr.DataArray | None:
    """Extract ERA5 features from the model state.

    Feature naming patterns:
    - Without aggregations: `era5_{variable}_{time}`
      Example: `era5_temperature_2020_summer`
    - With aggregations: `era5_{variable}_{time}_{agg}`
      Example: `era5_temperature_2020_summer_mean`

    Args:
        model_state: The xarray Dataset containing the model state.
        importance_type: Type of feature importance to extract.
        temporal_group: Filter to specific temporal group ('yearly', 'seasonal', 'shoulder').
                       If None, extracts all ERA5 features.

    Returns:
        xr.DataArray: The extracted ERA5 features. This DataArray has dimensions
            ('variable', 'time') or ('variable', 'time', 'agg') depending on the feature structure.
            Returns None if no ERA5 features are found.

    """

    def _is_era5_feature(feature: str) -> bool:
        return feature.startswith("era5_")

    def _get_temporal_group(feature: str) -> str:
        """Determine temporal group from feature name based on time component.

        Time patterns:
        - Yearly: just a year number (e.g., "2020")
        - Seasonal: season_year (e.g., "summer_2020", "winter_2021")
        - Shoulder: SHOULDER_year (e.g., "JFM_2020", "OND_2021")
        """
        parts = feature.split("_")
        common_aggs = {
            "mean",
            "std",
            "min",
            "max",
            "median",
            "sum",
            "count",
            "q25",
            "q75",
        }

        # Find where the time part starts (after "era5" and variable name)
        # Pattern: era5_variable_time or era5_variable_time_agg
        # Time can be: year, season_year, or SHOULDER_year

        # Remove era5 prefix and aggregation suffix if present
        if parts[-1] in common_aggs:
            time_parts = parts[1:-1]  # Remove era5 and agg
        else:
            time_parts = parts[1:]  # Remove era5

        # The last part (or last two parts) is the time
        # Check last part first
        last_part = time_parts[-1]

        if last_part.isdigit():
            # Could be yearly (just year) or part of seasonal/shoulder (season_year or SHOULDER_year)
            if len(time_parts) >= 2:
                second_last = time_parts[-2]
                # Check if second to last is a season or shoulder indicator
                if second_last.lower() in ["summer", "winter"]:
                    return "seasonal"
                elif second_last.upper() in ["JFM", "AMJ", "JAS", "OND"]:
                    return "shoulder"
            # Just a year number
            return "yearly"
        elif last_part.lower() in ["summer", "winter"]:
            return "seasonal"
        elif last_part.upper() in ["JFM", "AMJ", "JAS", "OND"]:
            return "shoulder"

        return "unknown"

    def _extract_var_name(feature: str) -> str:
        parts = feature.split("_")
        # Pattern: era5_variable_year_agg OR era5_variable_season_year_agg
        # Variable can include stat names like t2m_max, t2m_min, etc.
        # We need to find where the year starts (it's always a 4-digit number)

        # Find the first year (4-digit number)
        year_idx = None
        for i, part in enumerate(parts):
            if part.isdigit() and len(part) == 4:
                year_idx = i
                break

        if year_idx is None:
            # Fallback: return everything after era5
            return "_".join(parts[1:])

        # Variable is everything between era5 and year
        # But also check if there's a season before the year
        seasons = {"summer", "winter"}
        shoulders = {"JFM", "AMJ", "JAS", "OND"}

        var_end_idx = year_idx
        # Check if the part before year is a season/shoulder
        if year_idx > 1 and (parts[year_idx - 1].lower() in seasons or parts[year_idx - 1].upper() in shoulders):
            var_end_idx = year_idx - 1

        # Variable is from index 1 (after era5) to var_end_idx
        return "_".join(parts[1:var_end_idx])

    def _extract_time_name(feature: str) -> str:
        parts = feature.split("_")
        common_aggs = {
            "mean",
            "std",
            "min",
            "max",
            "median",
            "sum",
            "count",
            "q25",
            "q75",
        }
        if parts[-1] in common_aggs:
            # Has aggregation: era5_var_time_agg -> time is second to last
            return parts[-2]
        else:
            # No aggregation: era5_var_time -> time is last
            return parts[-1]

    def _extract_season(feature: str) -> str | None:
        """Extract season/shoulder from seasonal/shoulder features.

        Pattern: era5_variable_season_year_agg or era5_variable_season_year
        """
        # Look through parts to find season/shoulder indicators
        for part in feature.split("_"):
            if part.lower() in ["summer", "winter"]:
                return part.lower()
            elif part.upper() in ["JFM", "AMJ", "JAS", "OND"]:
                return part.upper()

        return None

    def _extract_year_from_time(feature: str) -> str:
        """Extract just the year component from time.

        For seasonal/shoulder features, find the year that comes after the season.
        """
        parts = feature.split("_")
        common_aggs = {
            "mean",
            "std",
            "min",
            "max",
            "median",
            "sum",
            "count",
            "q25",
            "q75",
        }

        # Find the season/shoulder part, then the next part should be the year
        for i, part in enumerate(parts):
            if part.lower() in ["summer", "winter"] or part.upper() in [
                "JFM",
                "AMJ",
                "JAS",
                "OND",
            ]:
                # Next part should be the year
                if i + 1 < len(parts):
                    next_part = parts[i + 1]
                    # Skip if it's an aggregation, year should be before agg
                    if next_part not in common_aggs:
                        return next_part

        # Fallback: return last numeric part that's not an aggregation
        for part in reversed(parts):
            if part.isdigit():
                return part

        return "unknown"

    def _extract_agg_name(feature: str) -> str | None:
        parts = feature.split("_")
        common_aggs = {
            "mean",
            "std",
            "min",
            "max",
            "median",
            "sum",
            "count",
            "q25",
            "q75",
        }
        if parts[-1] in common_aggs:
            return parts[-1]
        return None

    # Get all ERA5 features
    era5_features = [f for f in model_state.feature.to_numpy() if _is_era5_feature(f)]

    # Filter by temporal group if specified
    if temporal_group is not None:
        era5_features = [f for f in era5_features if _get_temporal_group(f) == temporal_group]

    if len(era5_features) == 0:
        return None

    # Get the appropriate importance array
    importance_array = _get_feature_importance_array(model_state, importance_type)

    # Check if features have aggregations
    has_agg = any(_extract_agg_name(f) is not None for f in era5_features)

    # Check if features have season/shoulder split (for seasonal/shoulder temporal groups)
    has_season = any(_extract_season(f) is not None for f in era5_features)

    # Split the single feature dimension of era5 features into separate dimensions
    era5_features_array = importance_array.sel(feature=era5_features)

    if has_season:
        # For seasonal/shoulder: split into variable, season, year, (agg)
        era5_features_array = era5_features_array.assign_coords(
            variable=("feature", [_extract_var_name(f) for f in era5_features]),
            season=("feature", [_extract_season(f) for f in era5_features]),
            year=("feature", [_extract_year_from_time(f) for f in era5_features]),
        )

        if has_agg:
            era5_features_array = era5_features_array.assign_coords(
                agg=(
                    "feature",
                    [_extract_agg_name(f) or "none" for f in era5_features],
                ),
            )
            era5_features_array = era5_features_array.set_index(feature=["variable", "season", "year", "agg"]).unstack(
                "feature"
            )
        else:
            era5_features_array = era5_features_array.set_index(feature=["variable", "season", "year"]).unstack(
                "feature"
            )
    else:
        # For yearly: keep as variable, time, (agg)
        era5_features_array = era5_features_array.assign_coords(
            variable=("feature", [_extract_var_name(f) for f in era5_features]),
            time=("feature", [_extract_time_name(f) for f in era5_features]),
        )

        if has_agg:
            # Add aggregation dimension
            era5_features_array = era5_features_array.assign_coords(
                agg=(
                    "feature",
                    [_extract_agg_name(f) or "none" for f in era5_features],
                ),
            )
            era5_features_array = era5_features_array.set_index(feature=["variable", "time", "agg"]).unstack("feature")
        else:
            era5_features_array = era5_features_array.set_index(feature=["variable", "time"]).unstack("feature")

    return era5_features_array


def extract_arcticdem_features(
    model_state: xr.Dataset, importance_type: str = "feature_weights"
) -> xr.DataArray | None:
    """Extract ArcticDEM features from the model state.

    Feature naming pattern: `arcticdem_{variable}_{agg}`
    Example: `arcticdem_elevation_mean`, `arcticdem_slope_std`

    Args:
        model_state: The xarray Dataset containing the model state.
        importance_type: Type of feature importance to extract.

    Returns:
        xr.DataArray: The extracted ArcticDEM features. This DataArray has dimensions
            ('variable', 'agg') corresponding to the different components of the ArcticDEM features.
            Returns None if no ArcticDEM features are found.

    """

    def _is_arcticdem_feature(feature: str) -> bool:
        return feature.startswith("arcticdem_")

    def _extract_var_name(feature: str) -> str:
        parts = feature.split("_")
        # arcticdem_variable_agg format
        # Variable name is everything between arcticdem_ and the last part
        return "_".join(parts[1:-1])

    def _extract_agg_name(feature: str) -> str:
        parts = feature.split("_")
        # Last part is the aggregation
        return parts[-1]

    arcticdem_features = [f for f in model_state.feature.to_numpy() if _is_arcticdem_feature(f)]
    if len(arcticdem_features) == 0:
        return None

    # Get the appropriate importance array
    importance_array = _get_feature_importance_array(model_state, importance_type)

    # Split the single feature dimension into separate dimensions (variable, agg)
    arcticdem_feature_array = importance_array.sel(feature=arcticdem_features)
    arcticdem_feature_array = arcticdem_feature_array.assign_coords(
        variable=("feature", [_extract_var_name(f) for f in arcticdem_features]),
        agg=("feature", [_extract_agg_name(f) for f in arcticdem_features]),
    )
    arcticdem_feature_array = arcticdem_feature_array.set_index(feature=["variable", "agg"]).unstack("feature")
    return arcticdem_feature_array


def extract_common_features(model_state: xr.Dataset, importance_type: str = "feature_weights") -> xr.DataArray | None:
    """Extract common features (cell_area, water_area, land_area, land_ratio, lon, lat) from the model state.

    Args:
        model_state: The xarray Dataset containing the model state.
        importance_type: Type of feature importance to extract.

    Returns:
        xr.DataArray: The extracted common features with a single 'feature' dimension.
            Returns None if no common features are found.

    """
    common_feature_names = [
        "cell_area",
        "water_area",
        "land_area",
        "land_ratio",
        "lon",
        "lat",
    ]

    def _is_common_feature(feature: str) -> bool:
        return feature in common_feature_names

    common_features = [f for f in model_state.feature.to_numpy() if _is_common_feature(f)]
    if len(common_features) == 0:
        return None

    # Get the appropriate importance array
    importance_array = _get_feature_importance_array(model_state, importance_type)

    # Extract the feature importance for common features
    common_feature_array = importance_array.sel(feature=common_features)
    return common_feature_array
