# ruff: noqa: N806
"""Inference runs on trained models."""

from typing import Literal

import cupy as cp
import geopandas as gpd
import numpy as np
import pandas as pd
import torch
from autogluon.tabular import TabularPredictor
from rich import pretty, traceback
from sklearn import set_config
from sklearn.pipeline import Pipeline

from entropice.ml.dataset import DatasetEnsemble
from entropice.ml.models import SupportedModel, is_classifier, is_regressor
from entropice.utils.types import Task

traceback.install()
pretty.install()

set_config(array_api_dispatch=True)


def _assert_right_model_for_task(model: Pipeline | SupportedModel | TabularPredictor, task: Task) -> None:
    assert hasattr(model, "predict") or hasattr(model, "predict_proba"), (
        "Model must have a predict or predict_proba method"
    )
    is_classification = task in ["binary", "count_regimes", "density_regimes"]
    is_regression = task in ["count", "density"]
    assert is_classification != is_regression, "Task must be either classification or regression"
    assert is_classifier(model) == is_classification, f"Model type does not match task type for task {task}"
    assert is_regressor(model) == is_regression, f"Model type does not match task type for task {task}"


def _categorize_predictions(
    preds: pd.Series | np.ndarray,
    task: Task,
) -> pd.Series | pd.Categorical:
    """Convert the raw model predictions into a category type series."""
    if isinstance(preds, np.ndarray):
        preds = pd.Series(preds)
    match task:
        case "binary" | "count_regimes" | "density_regimes":
            labels_dict = {
                "binary": ["No RTS", "RTS"],
                "count_regimes": ["None", "Very Few", "Few", "Several", "Many", "Very Many"],
                "density_regimes": ["Empty", "Very Sparse", "Sparse", "Moderate", "Dense", "Very Dense"],
            }
            categories = pd.CategoricalDtype(categories=labels_dict[task], ordered=task != "binary")
            # Check if preds are codes or labels
            if preds.dtype == "object" or isinstance(preds.iloc[0], str):
                return pd.Categorical(preds, dtype=categories)
            else:
                return pd.Categorical.from_codes(preds.astype(int).to_list(), dtype=categories)
        case _:
            return preds


def predict_proba(
    e: DatasetEnsemble,
    model: Pipeline | SupportedModel | TabularPredictor,
    task: Task,
    device: Literal["cpu", "cuda", "torch"] = "cuda",
) -> gpd.GeoDataFrame:
    """Get predicted probabilities for each cell.

    Args:
        e (DatasetEnsemble): The dataset ensemble configuration.
        model: SupportedModel | TabularPredictor: The trained model to use for predictions.
        task (Task): The task.
        device (Literal["cpu", "cuda", "torch"]): The device to use for predictions.
            This must match with the state of the model!

    Returns:
        gpd.GeoDataFrame: A GeoDataFrame with cell_id, predicted probability, and geometry.

    """
    _assert_right_model_for_task(model, task)
    # Predict in batches to avoid memory issues
    batch_size = 50000
    preds = []

    grid_gdf = e.read_grid()
    for batch in e.create_inference_df(batch_size=batch_size):
        # Filter rows containing NaN values
        batch = batch.dropna(axis="index", how="any")

        # Skip empty batches (all rows had NaN values)
        if len(batch) == 0:
            continue

        cell_ids = batch.index.to_numpy()
        cell_geoms = grid_gdf.loc[batch.index, "geometry"].to_numpy()

        if isinstance(model, TabularPredictor):
            print(f"Predicting batch of size {len(batch)} ({type(batch)}) with AutoGluon TabularPredictor...")
            try:
                batch_preds = model.predict(batch)
            except Exception as ex:
                print("Something went wrong")
                print(batch)
                print(batch.columns)
                raise ex
            print(f"Batch predictions type: {type(batch_preds)}, shape: {batch_preds.shape}")

            assert isinstance(batch_preds, pd.DataFrame | pd.Series), (
                "AutoGluon predict should return a DataFrame or Series"
            )
            batch_preds = batch_preds.to_numpy()
        else:
            X_batch = batch.to_numpy(dtype="float64")
            if device == "torch":
                X_batch = torch.from_numpy(X_batch).to("cuda")
            elif device == "cuda":
                X_batch = cp.asarray(X_batch)
            batch_preds = model.predict(X_batch)
            if isinstance(batch_preds, cp.ndarray):
                batch_preds = batch_preds.get()
            elif torch.is_tensor(batch_preds):
                batch_preds = batch_preds.cpu().numpy()
            batch_preds = _categorize_predictions(batch_preds, task=task)
        batch_preds = gpd.GeoDataFrame(
            {
                "cell_id": cell_ids,
                "predicted": batch_preds,
                "geometry": cell_geoms,
            },
        ).set_crs(grid_gdf.crs, inplace=False)
        preds.append(batch_preds)
    return gpd.GeoDataFrame(pd.concat(preds, ignore_index=True))
