"""Inference scrips."""

from __future__ import annotations

from functools import lru_cache
from typing import Any

import pandas as pd

from vito_crop_classification.filters import BaseFilter
from vito_crop_classification.inference.patch import smooth_patch
from vito_crop_classification.inference.predictor import Predictor
from vito_crop_classification.model import Model


@lru_cache(maxsize=1)
def load_predictor(
    model: Model,
    filter_cls: Any = BaseFilter,
) -> Predictor:
    """
    Load in a model used for predictions.

    Parameters
    ----------
    model : Model, optional
        Model to use for inference
    filter_cls: Any
        Filter class used to create the filter applied during inference

    Returns
    -------
    predictor : Predictor
        Predictor class to run inference with

    Note
    ----
     - Either model or mdl_f should be provided
    """
    return Predictor(
        model=model,
        filterer=filter_cls(classes=model.get_class_ids(), filter_f=model.model_folder / "filters"),
    )


def main(
    df: pd.DataFrame,
    model: Model,
    patch_smoothing: bool = False,
    patch_shape: tuple[int, int] | None = None,
    transform: bool = True,
    ignore_classes: list[str] = [],
) -> pd.DataFrame:
    """
    Run the inference pipeline.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame to make predictions on
    model : Model, optional
        Model to use for inference
    patch_smoothing : bool
        Smooth on a patch level
    patch_shape : tuple[int, int], optional
        Shape of the patch, should be provided if patch_smoothing=True
    transform : bool, optional
        Run dataframe transformation
    ignore_classes : list[str], optional
        Define classes to ignore while predicting

    Returns
    -------
    prediction : pd.DataFrame
        Prediction on the provided dataframe

    Note
    ----
     - Either model or mdl_f should be provided
    """
    # Make pixel-level predictions
    predictor = load_predictor(model=model)
    pred_df = predictor(df, transform=transform, ignore_classes=ignore_classes)

    # Apply patch smoothing if requested
    if patch_smoothing:
        smooth_patch(
            df=pred_df,
            shape=patch_shape,
            class_ids=predictor.model.get_class_ids(),
            class_names=predictor.model.get_class_names(),
        )

    # Return the final result
    return pred_df
