from __future__ import annotations

from contextlib import contextmanager
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    Iterable,
    Literal,
    Sequence,
    cast,
    overload,
)

from tqdm.auto import trange

from ._utils.common import UNSET, CreateMode, DropMode, logger
from .client import (
    ListPredictionsRequest,
    OrcaClient,
    PredictiveModelUpdate,
    RARHeadType,
)
from .client import RegressionMetrics as RegressionMetricsResponse
from .client import (
    RegressionModelMetadata,
    RegressionPredictionRequest,
)
from .datasource import Datasource
from .job import Job
from .memoryset import ConsistencyLevel, ScoredMemoryset
from .telemetry import (
    RegressionPrediction,
    TelemetryMode,
    _get_telemetry_config,
    _parse_feedback,
)

if TYPE_CHECKING:
    # Peer dependency - user has datasets if they have a Dataset object
    from datasets import Dataset as HFDataset  # type: ignore
    from pandas import DataFrame as PandasDataFrame  # type: ignore


class RegressionMetrics:
    """
    Metrics for evaluating regression model performance.

    Attributes:
        coverage: Percentage of predictions that are not none
        mse: Mean squared error of the predictions
        rmse: Root mean squared error of the predictions
        mae: Mean absolute error of the predictions
        r2: R-squared score (coefficient of determination) of the predictions
        explained_variance: Explained variance score of the predictions
        loss: Mean squared error loss of the predictions
        anomaly_score_mean: Mean of anomaly scores across the dataset
        anomaly_score_median: Median of anomaly scores across the dataset
        anomaly_score_variance: Variance of anomaly scores across the dataset
    """

    coverage: float
    mse: float
    rmse: float
    mae: float
    r2: float
    explained_variance: float
    loss: float
    anomaly_score_mean: float | None
    anomaly_score_median: float | None
    anomaly_score_variance: float | None
    model: RegressionModel | None

    def __init__(self, response: RegressionMetricsResponse, model: RegressionModel | None = None):
        self.coverage = response["coverage"]
        self.mse = response["mse"]
        self.rmse = response["rmse"]
        self.mae = response["mae"]
        self.r2 = response["r2"]
        self.explained_variance = response["explained_variance"]
        self.loss = response["loss"]
        self.anomaly_score_mean = response.get("anomaly_score_mean")
        self.anomaly_score_median = response.get("anomaly_score_median")
        self.anomaly_score_variance = response.get("anomaly_score_variance")
        self.model = model
        for warning in response.get("warnings", []):
            logger.warning(warning)

    def __repr__(self) -> str:
        return (
            "RegressionMetrics({\n"
            + f"    mae: {self.mae:.4f},\n"
            + f"    rmse: {self.rmse:.4f},\n"
            + f"    r2: {self.r2:.4f},\n"
            + (
                f"    anomaly_score: {self.anomaly_score_mean:.4f} ± {self.anomaly_score_variance:.4f},\n"
                if self.anomaly_score_mean
                else ""
            )
            + "})"
        )

    @classmethod
    def compute(
        cls,
        predictions: Sequence[RegressionPrediction],
    ) -> RegressionMetrics:
        """
        Compute regression metrics from a list of predictions.

        Params:
            predictions: List of RegressionPrediction objects with expected_score set

        Returns:
            RegressionMetrics with computed metrics

        Raises:
            ValueError: If any prediction is missing expected_score
        """
        if len(predictions) > 100_000:
            raise ValueError("Too many predictions, maximum is 100,000")
        if len(predictions) == 0:
            raise ValueError("No predictions provided")
        if any(p.expected_score is None for p in predictions):
            raise ValueError("All predictions must have expected_score set")
        expected_scores = [cast(float, p.expected_score) for p in predictions]
        predicted_scores = [p.score for p in predictions]
        anomaly_scores = (
            None
            if any(p.anomaly_score is None for p in predictions)
            else [cast(float, p.anomaly_score) for p in predictions]
        )

        client = OrcaClient._resolve_client()
        response = client.POST(
            "/regression_model/metrics",
            json={
                "expected_scores": expected_scores,
                "predicted_scores": predicted_scores,
                "anomaly_scores": anomaly_scores,
            },
        )
        return cls(response, model=predictions[0].model)


class RegressionModel:
    """
    A handle to a regression model in OrcaCloud

    Attributes:
        id: Unique identifier for the model
        name: Unique name of the model
        description: Optional description of the model
        memoryset: Memoryset that the model uses
        head_type: Regression head type of the model
        memory_lookup_count: Number of memories the model uses for each prediction
        locked: Whether the model is locked to prevent accidental deletion
        created_at: When the model was created
        updated_at: When the model was last updated
    """

    id: str
    name: str
    description: str | None
    memoryset: ScoredMemoryset
    head_type: RARHeadType
    memory_lookup_count: int
    version: int
    locked: bool
    created_at: datetime
    updated_at: datetime
    memoryset_id: str

    _last_prediction: RegressionPrediction | None
    _last_prediction_was_batch: bool
    _memoryset_override: ScoredMemoryset | None

    def __init__(self, metadata: RegressionModelMetadata):
        # for internal use only, do not document
        self.id = metadata["id"]
        self.name = metadata["name"]
        self.description = metadata["description"]
        self.memoryset = ScoredMemoryset.open(metadata["memoryset_id"])
        self.head_type = metadata["head_type"]
        self.memory_lookup_count = metadata["memory_lookup_count"]
        self.version = metadata["version"]
        self.locked = metadata["locked"]
        self.created_at = datetime.fromisoformat(metadata["created_at"])
        self.updated_at = datetime.fromisoformat(metadata["updated_at"])
        self.memoryset_id = metadata["memoryset_id"]

        self._memoryset_override = None
        self._last_prediction = None
        self._last_prediction_was_batch = False

    def __eq__(self, other) -> bool:
        return isinstance(other, RegressionModel) and self.id == other.id

    def __repr__(self):
        memoryset_repr = self.memoryset.__repr__().replace("\n", "\n    ")
        return (
            "RegressionModel({\n"
            f"    name: '{self.name}',\n"
            f"    memory_lookup_count: {self.memory_lookup_count},\n"
            f"    memoryset: {memoryset_repr},\n"
            "})"
        )

    @property
    def last_prediction(self) -> RegressionPrediction:
        """
        Last prediction made by the model

        Note:
            If the last prediction was part of a batch prediction, the last prediction from the
            batch is returned. If no prediction has been made yet, a [`LookupError`][LookupError]
            is raised.
        """
        if self._last_prediction_was_batch:
            logger.warning(
                "Last prediction was part of a batch prediction, returning the last prediction from the batch"
            )
        if self._last_prediction is None:
            raise LookupError("No prediction has been made yet")
        return self._last_prediction

    @classmethod
    def create(
        cls,
        name: str,
        memoryset: ScoredMemoryset,
        memory_lookup_count: int | None = None,
        description: str | None = None,
        if_exists: CreateMode | Literal["replace"] = "error",
    ) -> RegressionModel:
        """
        Create a regression model.

        Params:
            name: Name of the model
            memoryset: The scored memoryset to use for prediction
            memory_lookup_count: Number of memories to retrieve for prediction. Defaults to 10.
            description: Description of the model
            if_exists: What to do if a model with the same name already exists:
                - `"error"`: raise an error if the model already exists
                - `"open"`: open the existing model
                - `"replace"`: drop the existing model including all its prediction telemetry

        Returns:
            RegressionModel instance

        Raises:
            ValueError: If a model with the same name already exists and if_exists is "error"
            ValueError: If the memoryset is empty
            ValueError: If memory_lookup_count exceeds the number of memories in the memoryset
        """
        existing = cls.exists(name)
        if existing:
            if if_exists == "error":
                raise ValueError(f"RegressionModel with name '{name}' already exists")
            elif if_exists == "replace":
                cls.drop(name, if_not_exists="ignore")
            elif if_exists == "open":
                existing = cls.open(name)
                for attribute in {"memory_lookup_count"}:
                    local_attribute = locals()[attribute]
                    existing_attribute = getattr(existing, attribute)
                    if local_attribute is not None and local_attribute != existing_attribute:
                        raise ValueError(f"Model with name {name} already exists with different {attribute}")

                # special case for memoryset
                if existing.memoryset_id != memoryset.id:
                    raise ValueError(f"Model with name {name} already exists with different memoryset")

                return existing

        client = OrcaClient._resolve_client()
        metadata = client.POST(
            "/regression_model",
            json={
                "name": name,
                "memoryset_name_or_id": memoryset.id,
                "memory_lookup_count": memory_lookup_count,
                "description": description,
            },
        )
        return cls(metadata)

    @classmethod
    def open(cls, name: str) -> RegressionModel:
        """
        Get a handle to a regression model in the OrcaCloud

        Params:
            name: Name or unique identifier of the regression model

        Returns:
            Handle to the existing regression model in the OrcaCloud

        Raises:
            LookupError: If the regression model does not exist
        """
        client = OrcaClient._resolve_client()
        return cls(client.GET("/regression_model/{name_or_id}", params={"name_or_id": name}))

    @classmethod
    def exists(cls, name_or_id: str) -> bool:
        """
        Check if a regression model exists in the OrcaCloud

        Params:
            name_or_id: Name or id of the regression model

        Returns:
            `True` if the regression model exists, `False` otherwise
        """
        try:
            cls.open(name_or_id)
            return True
        except LookupError:
            return False

    @classmethod
    def all(cls) -> list[RegressionModel]:
        """
        Get a list of handles to all regression models in the OrcaCloud

        Returns:
            List of handles to all regression models in the OrcaCloud
        """
        client = OrcaClient._resolve_client()
        return [cls(metadata) for metadata in client.GET("/regression_model", params={})]

    @classmethod
    def drop(cls, name_or_id: str, if_not_exists: DropMode = "error"):
        """
        Delete a regression model from the OrcaCloud

        Warning:
            This will delete the model and all associated data, including predictions, evaluations, and feedback.

        Params:
            name_or_id: Name or id of the regression model
            if_not_exists: What to do if the regression model does not exist, defaults to `"error"`.
                Other option is `"ignore"` to do nothing if the regression model does not exist.

        Raises:
            LookupError: If the regression model does not exist and if_not_exists is `"error"`
        """
        try:
            client = OrcaClient._resolve_client()
            client.DELETE("/regression_model/{name_or_id}", params={"name_or_id": name_or_id})
            logger.info(f"Deleted model {name_or_id}")
        except LookupError:
            if if_not_exists == "error":
                raise

    def refresh(self):
        """Refresh the model data from the OrcaCloud"""
        self.__dict__.update(self.open(self.name).__dict__)

    def set(self, *, description: str | None = UNSET, locked: bool = UNSET) -> None:
        """
        Update editable attributes of the model.

        Note:
            If a field is not provided, it will default to [UNSET][orca_sdk.UNSET] and not be updated.

        Params:
            description: Value to set for the description
            locked: Value to set for the locked status

        Examples:
            Update the description:
            >>> model.set(description="New description")

            Remove description:
            >>> model.set(description=None)

            Lock the model:
            >>> model.set(locked=True)
        """
        update: PredictiveModelUpdate = {}
        if description is not UNSET:
            update["description"] = description
        if locked is not UNSET:
            update["locked"] = locked
        client = OrcaClient._resolve_client()
        client.PATCH("/regression_model/{name_or_id}", params={"name_or_id": self.id}, json=update)
        self.refresh()

    def lock(self) -> None:
        """Lock the model to prevent accidental deletion"""
        self.set(locked=True)

    def unlock(self) -> None:
        """Unlock the model to allow deletion"""
        self.set(locked=False)

    @overload
    def predict(
        self,
        value: str,
        expected_scores: float | None = None,
        tags: set[str] | None = None,
        save_telemetry: TelemetryMode = "on",
        prompt: str | None = None,
        use_lookup_cache: bool = True,
        timeout_seconds: int = 10,
        ignore_unlabeled: bool = False,
        partition_id: str | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
        use_gpu: bool = True,
        batch_size: int = 100,
        consistency_level: ConsistencyLevel = "Bounded",
    ) -> RegressionPrediction: ...

    @overload
    def predict(
        self,
        value: list[str],
        expected_scores: list[float] | None = None,
        tags: set[str] | None = None,
        save_telemetry: TelemetryMode = "on",
        prompt: str | None = None,
        use_lookup_cache: bool = True,
        timeout_seconds: int = 10,
        ignore_unlabeled: bool = False,
        partition_id: str | list[str | None] | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
        use_gpu: bool = True,
        batch_size: int = 100,
        consistency_level: ConsistencyLevel = "Bounded",
    ) -> list[RegressionPrediction]: ...

    # TODO: add filter support
    def predict(
        self,
        value: str | list[str],
        expected_scores: float | list[float] | None = None,
        tags: set[str] | None = None,
        save_telemetry: TelemetryMode = "on",
        prompt: str | None = None,
        use_lookup_cache: bool = True,
        timeout_seconds: int = 10,
        ignore_unlabeled: bool = False,
        partition_id: str | list[str | None] | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
        use_gpu: bool = True,
        batch_size: int = 100,
        consistency_level: ConsistencyLevel = "Bounded",
    ) -> RegressionPrediction | list[RegressionPrediction]:
        """
        Make predictions using the regression model.

        Params:
            value: Input text(s) to predict scores for
            expected_scores: Expected score(s) for telemetry tracking
            tags: Tags to associate with the prediction(s)
            save_telemetry: Whether to save telemetry for the prediction(s), defaults to `True`,
                which will save telemetry asynchronously unless the `ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY`
                environment variable is set to `"1"`. You can also pass `"sync"` or `"async"` to
                explicitly set the save mode.
            prompt: Optional prompt for instruction-tuned embedding models
            use_lookup_cache: Whether to use cached lookup results for faster predictions
            timeout_seconds: Timeout in seconds for the request, defaults to 10 seconds
            ignore_unlabeled: If True, only use memories with scores during lookup.
                If False (default), allow memories without scores when necessary.
            partition_id: Optional partition ID(s) to use during memory lookup
            partition_filter_mode: Optional partition filter mode to use for the prediction(s). One of
                * `"ignore_partitions"`: Ignore partitions
                * `"include_global"`: Include global memories
                * `"exclude_global"`: Exclude global memories
                * `"only_global"`: Only include global memories
            use_gpu: Whether to use GPU for the prediction (defaults to True)
            batch_size: Number of values to process in a single API call

        Returns:
            Single RegressionPrediction or list of RegressionPrediction objects

        Raises:
            ValueError: If expected_scores length doesn't match value length for batch predictions
            ValueError: If timeout_seconds is not a positive integer
            TimeoutError: If the request times out after the specified duration
        """
        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds must be a positive integer")
        if batch_size <= 0 or batch_size > 500:
            raise ValueError("batch_size must be between 1 and 500")

        if use_gpu:
            endpoint = "/gpu/regression_model/{name_or_id}/prediction"
        else:
            endpoint = "/regression_model/{name_or_id}/prediction"

        telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
        client = OrcaClient._resolve_client()

        # Convert to list for batching
        values = [value] if isinstance(value, str) else list(value)
        if isinstance(expected_scores, list) and len(expected_scores) != len(values):
            raise ValueError("Invalid input: \n\texpected_scores must be the same length as values")
        if isinstance(partition_id, list) and len(partition_id) != len(values):
            raise ValueError("Invalid input: \n\tpartition_id must be the same length as values")

        if expected_scores is not None and isinstance(expected_scores, (float, int)):
            expected_scores = [float(expected_scores)] * len(values)

        predictions: list[RegressionPrediction] = []
        for i in trange(0, len(values), batch_size, desc="predict", disable=len(values) <= batch_size):
            batch_values = values[i : i + batch_size]
            batch_expected_scores = expected_scores[i : i + batch_size] if expected_scores else None

            request_json: RegressionPredictionRequest = {
                "input_values": batch_values,
                "memoryset_override_name_or_id": self._memoryset_override.id if self._memoryset_override else None,
                "expected_scores": batch_expected_scores,
                "tags": list(tags or set()),
                "save_telemetry": telemetry_on,
                "save_telemetry_synchronously": telemetry_sync,
                "prompt": prompt,
                "use_lookup_cache": use_lookup_cache,
                "ignore_unlabeled": ignore_unlabeled,
                "partition_filter_mode": partition_filter_mode,
                "consistency_level": consistency_level,
            }
            if partition_filter_mode != "ignore_partitions":
                request_json["partition_ids"] = (
                    partition_id[i : i + batch_size] if isinstance(partition_id, list) else partition_id
                )

            response = client.POST(
                endpoint,
                params={"name_or_id": self.id},
                json=request_json,
                timeout=timeout_seconds,
            )

            if telemetry_on and any(p["prediction_id"] is None for p in response):
                raise RuntimeError("Failed to save prediction to database.")

            batch_expected = batch_expected_scores or [None] * len(batch_values)
            predictions.extend(
                RegressionPrediction(
                    prediction_id=prediction["prediction_id"],
                    label=None,
                    label_name=None,
                    score=prediction["score"],
                    confidence=prediction["confidence"],
                    anomaly_score=prediction["anomaly_score"],
                    memoryset=self._memoryset_override if self._memoryset_override is not None else self.memoryset,
                    model=self,
                    logits=None,
                    input_value=input_value,
                    expected_score=exp_score,
                )
                for prediction, input_value, exp_score in zip(response, batch_values, batch_expected)
            )

        self._last_prediction_was_batch = isinstance(value, list)
        self._last_prediction = predictions[-1]
        return predictions if isinstance(value, list) else predictions[0]

    def predictions(
        self,
        limit: int | None = None,
        offset: int = 0,
        tag: str | None = None,
        sort: list[tuple[Literal["anomaly_score", "confidence", "timestamp"], Literal["asc", "desc"]]] = [],
        batch_size: int = 100,
    ) -> list[RegressionPrediction]:
        """
        Get a list of predictions made by this model

        Params:
            limit: Maximum number of predictions to return. If `None`, returns all predictions
                by automatically paginating through results.
            offset: Optional offset of the first prediction to return
            tag: Optional tag to filter predictions by
            sort: Optional list of columns and directions to sort the predictions by.
                Predictions can be sorted by `created_at`, `confidence`, `anomaly_score`, or `score`.
            batch_size: Number of predictions to fetch in a single API call

        Returns:
            List of score predictions

        Examples:
            Get all predictions with a specific tag:
            >>> predictions = model.predictions(tag="evaluation")

            Get the last 3 predictions:
            >>> predictions = model.predictions(limit=3, sort=[("created_at", "desc")])
            [
                RegressionPrediction({score: 4.5, confidence: 0.95, anomaly_score: 0.1, input_value: 'Great service'}),
                RegressionPrediction({score: 2.0, confidence: 0.90, anomaly_score: 0.1, input_value: 'Poor experience'}),
                RegressionPrediction({score: 3.5, confidence: 0.85, anomaly_score: 0.1, input_value: 'Average'}),
            ]

            Get second most confident prediction:
            >>> predictions = model.predictions(sort=[("confidence", "desc")], offset=1, limit=1)
            [RegressionPrediction({score: 4.2, confidence: 0.90, anomaly_score: 0.1, input_value: 'Good service'})]
        """
        if batch_size <= 0 or batch_size > 500:
            raise ValueError("batch_size must be between 1 and 500")
        if limit == 0:
            return []

        client = OrcaClient._resolve_client()
        all_predictions: list[RegressionPrediction] = []

        # Cache memorysets to avoid repeated lookups
        memorysets: dict[str, ScoredMemoryset] = {self.memoryset_id: self.memoryset}

        def resolve_memoryset(memoryset_id: str) -> ScoredMemoryset:
            if memoryset_id not in memorysets:
                memorysets[memoryset_id] = ScoredMemoryset.open(memoryset_id)
            return memorysets[memoryset_id]

        if limit is not None and limit < batch_size:
            pages = [(offset, limit)]
        else:
            # automatically paginate the requests if necessary
            total = client.POST(
                "/telemetry/prediction/count",
                json={
                    "model_id": self.id,
                    "tag": tag,
                },
            )
            max_limit = max(total - offset, 0)
            limit = min(limit, max_limit) if limit is not None else max_limit
            pages = [(o, min(batch_size, limit - (o - offset))) for o in range(offset, offset + limit, batch_size)]

        for current_offset, current_limit in pages:
            request_json: ListPredictionsRequest = {
                "model_id": self.id,
                "limit": current_limit,
                "offset": current_offset,
                "tag": tag,
            }
            if sort:
                request_json["sort"] = sort
            response = client.POST(
                "/telemetry/prediction",
                json=request_json,
            )
            all_predictions.extend(
                RegressionPrediction(
                    prediction_id=prediction["prediction_id"],
                    label=None,
                    label_name=None,
                    score=prediction["score"],
                    confidence=prediction["confidence"],
                    anomaly_score=prediction["anomaly_score"],
                    memoryset=resolve_memoryset(prediction["memoryset_id"]),
                    model=self,
                    telemetry=prediction,
                    logits=None,
                    input_value=None,
                )
                for prediction in response
                if "score" in prediction
            )

        return all_predictions

    def _evaluate_datasource(
        self,
        datasource: Datasource,
        value_column: str,
        score_column: str,
        record_predictions: bool,
        tags: set[str] | None,
        subsample: int | float | None,
        background: bool = False,
        ignore_unlabeled: bool = False,
        partition_column: str | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
    ) -> RegressionMetrics | Job[RegressionMetrics]:
        client = OrcaClient._resolve_client()
        response = client.POST(
            "/regression_model/{model_name_or_id}/evaluation",
            params={"model_name_or_id": self.id},
            json={
                "datasource_name_or_id": datasource.id,
                "datasource_score_column": score_column,
                "datasource_value_column": value_column,
                "memoryset_override_name_or_id": self._memoryset_override.id if self._memoryset_override else None,
                "record_telemetry": record_predictions,
                "telemetry_tags": list(tags) if tags else None,
                "subsample": subsample,
                "ignore_unlabeled": ignore_unlabeled,
                "datasource_partition_column": partition_column,
                "partition_filter_mode": partition_filter_mode,
            },
        )

        def get_value():
            client = OrcaClient._resolve_client()
            res = client.GET(
                "/regression_model/{model_name_or_id}/evaluation/{job_id}",
                params={"model_name_or_id": self.id, "job_id": response["job_id"]},
            )
            assert res["result"] is not None
            return RegressionMetrics(res["result"], model=self)

        job = Job(response["job_id"], get_value)
        return job if background else job.result()

    def _evaluate_local(
        self,
        data: Iterable[dict[str, Any]],
        value_column: str,
        score_column: str,
        record_predictions: bool,
        tags: set[str],
        batch_size: int,
        prompt: str | None = None,
        ignore_unlabeled: bool = False,
        partition_column: str | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
    ) -> RegressionMetrics:
        values: list[str] = []
        expected_scores: list[float] = []
        partition_ids: list[str | None] | None = [] if partition_column else None

        for sample in data:
            if len(values) >= 100_000:
                raise ValueError("Upload a Datasource to evaluate against more than 100,000 samples.")
            values.append(sample[value_column])
            expected_score = sample[score_column]
            if expected_score is None:
                raise ValueError("Expected score is required for all samples")
            expected_scores.append(expected_score)
            if partition_ids is not None and partition_column:
                partition_ids.append(sample[partition_column])

        if not values:
            raise ValueError("Evaluation data cannot be empty")

        predictions = self.predict(
            values,
            expected_scores=expected_scores,
            tags=tags,
            save_telemetry="sync" if record_predictions else "off",
            prompt=prompt,
            ignore_unlabeled=ignore_unlabeled,
            partition_id=partition_ids,
            partition_filter_mode=partition_filter_mode,
            batch_size=batch_size,
        )

        return RegressionMetrics.compute(predictions)

    @overload
    def evaluate(
        self,
        data: Datasource,
        *,
        value_column: str = "value",
        score_column: str = "score",
        record_predictions: bool = False,
        tags: set[str] = {"evaluation"},
        batch_size: int = 100,
        prompt: str | None = None,
        subsample: int | float | None = None,
        background: Literal[True],
        ignore_unlabeled: bool = False,
        partition_column: str | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
    ) -> Job[RegressionMetrics]:
        pass

    @overload
    def evaluate(
        self,
        data: Datasource | HFDataset | PandasDataFrame | Iterable[dict[str, Any]],
        *,
        value_column: str = "value",
        score_column: str = "score",
        record_predictions: bool = False,
        tags: set[str] = {"evaluation"},
        batch_size: int = 100,
        prompt: str | None = None,
        subsample: int | float | None = None,
        background: Literal[False] = False,
        ignore_unlabeled: bool = False,
        partition_column: str | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
    ) -> RegressionMetrics:
        pass

    def evaluate(
        self,
        data: Datasource | HFDataset | PandasDataFrame | Iterable[dict[str, Any]],
        *,
        value_column: str = "value",
        score_column: str = "score",
        record_predictions: bool = False,
        tags: set[str] = {"evaluation"},
        batch_size: int = 100,
        prompt: str | None = None,
        subsample: int | float | None = None,
        background: bool = False,
        ignore_unlabeled: bool = False,
        partition_column: str | None = None,
        partition_filter_mode: Literal[
            "ignore_partitions", "include_global", "exclude_global", "only_global"
        ] = "include_global",
    ) -> RegressionMetrics | Job[RegressionMetrics]:
        """
        Evaluate the regression model on a given dataset or datasource

        Params:
            data: the data to evaluate the model on. This can be an Orca [`Datasource`][orca_sdk.datasource.Datasource],
                a Hugging Face [`Dataset`][datasets.Dataset], a pandas [`DataFrame`][pandas.DataFrame], or an iterable of dictionaries.
            value_column: Name of the column that contains the input values to the model
            score_column: Name of the column containing the expected scores
            record_predictions: Whether to record [`RegressionPrediction`][orca_sdk.telemetry.RegressionPrediction]s for analysis
            tags: Optional tags to add to the recorded [`RegressionPrediction`][orca_sdk.telemetry.RegressionPrediction]s
            batch_size: Batch size for processing the data inputs (not used for Datasource inputs)
            prompt: Optional prompt for instruction-tuned embedding models
            subsample: Optional number (int) of rows to sample or fraction (float in (0, 1]) of data to sample for evaluation.
            background: Whether to run the operation in the background and return a job handle
            ignore_unlabeled: If True, only use memories with scores during lookup. If False (default), allow memories without scores
            partition_column: Optional name of the column that contains the partition IDs
            partition_filter_mode: Optional partition filter mode to use for the evaluation. One of
                * `"ignore_partitions"`: Ignore partitions
                * `"include_global"`: Include global memories
                * `"exclude_global"`: Exclude global memories
                * `"only_global"`: Only include global memories
        Returns:
            RegressionMetrics containing metrics including MAE, MSE, RMSE, R2, and anomaly score statistics

        Examples:
            >>> model.evaluate(datasource, value_column="text", score_column="rating")
            RegressionMetrics({
                mae: 0.2500,
                rmse: 0.3536,
                r2: 0.8500,
                anomaly_score: 0.3500 ± 0.0500,
            })

            >>> # Using with an instruction-tuned embedding model
            >>> model.evaluate(dataset,prompt="Represent this review for rating prediction:")
            RegressionMetrics({
            mae: 0.2000,
            rmse: 0.3000,
            r2: 0.9000,
            anomaly_score: 0.3000 ± 0.0400})
        """
        if isinstance(data, Datasource):
            return self._evaluate_datasource(
                datasource=data,
                value_column=value_column,
                score_column=score_column,
                record_predictions=record_predictions,
                tags=tags,
                subsample=subsample,
                background=background,
                ignore_unlabeled=ignore_unlabeled,
                partition_column=partition_column,
                partition_filter_mode=partition_filter_mode,
            )
        else:
            if background:
                raise ValueError("Background evaluation is only supported for Datasource inputs")
            try:
                import pandas as pd  # type: ignore

                if isinstance(data, pd.DataFrame):
                    data = data.to_dict(orient="records")  # type: ignore
            except ImportError:
                pass

            if not hasattr(data, "__iter__"):
                raise ValueError(
                    f"Invalid data type: {type(data).__name__}. "
                    "Expected Iterable[dict], HuggingFace Dataset, or pandas DataFrame."
                )

            return self._evaluate_local(
                data=cast(Iterable[dict[str, Any]], data),
                value_column=value_column,
                score_column=score_column,
                record_predictions=record_predictions,
                tags=tags,
                batch_size=batch_size,
                prompt=prompt,
                ignore_unlabeled=ignore_unlabeled,
                partition_column=partition_column,
                partition_filter_mode=partition_filter_mode,
            )

    @contextmanager
    def use_memoryset(self, memoryset_override: ScoredMemoryset) -> Generator[None, None, None]:
        """
        Temporarily override the memoryset used by the model for predictions

        Params:
            memoryset_override: Memoryset to override the default memoryset with

        Examples:
            >>> with model.use_memoryset(ScoredMemoryset.open("my_other_memoryset")):
            ...     predictions = model.predict("Rate your experience")
        """
        self._memoryset_override = memoryset_override
        try:
            yield
        finally:
            self._memoryset_override = None

    @overload
    def record_feedback(self, feedback: dict[str, Any]) -> None:
        pass

    @overload
    def record_feedback(self, feedback: Iterable[dict[str, Any]]) -> None:
        pass

    def record_feedback(self, feedback: Iterable[dict[str, Any]] | dict[str, Any]):
        """
        Record feedback for a list of predictions.

        We support recording feedback in several categories for each prediction. A
        [`FeedbackCategory`][orca_sdk.telemetry.FeedbackCategory] is created automatically,
        the first time feedback with a new name is recorded. Categories are global across models.
        The value type of the category is inferred from the first recorded value. Subsequent
        feedback for the same category must be of the same type.

        Params:
            feedback: Feedback to record, this should be dictionaries with the following keys:

                - `category`: Name of the category under which to record the feedback.
                - `value`: Feedback value to record, should be `True` for positive feedback and
                    `False` for negative feedback or a [`float`][float] between `-1.0` and `+1.0`
                    where negative values indicate negative feedback and positive values indicate
                    positive feedback.
                - `comment`: Optional comment to record with the feedback.

        Examples:
            Record whether predictions were accurate:
            >>> model.record_feedback({
            ...     "prediction": p.prediction_id,
            ...     "category": "accurate",
            ...     "value": abs(p.score - p.expected_score) < 0.5,
            ... } for p in predictions)

            Record star rating as normalized continuous score between `-1.0` and `+1.0`:
            >>> model.record_feedback({
            ...     "prediction": "123e4567-e89b-12d3-a456-426614174000",
            ...     "category": "rating",
            ...     "value": -0.5,
            ...     "comment": "2 stars"
            ... })

        Raises:
            ValueError: If the value does not match previous value types for the category, or is a
                [`float`][float] that is not between `-1.0` and `+1.0`.
        """
        client = OrcaClient._resolve_client()
        client.PUT(
            "/telemetry/prediction/feedback",
            json=[
                _parse_feedback(f) for f in (cast(list[dict], [feedback]) if isinstance(feedback, dict) else feedback)
            ],
        )
