#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from typing import Any
from typing import Optional
from typing import SupportsFloat
from typing import SupportsInt
from typing import TypeVar

import apache_beam as beam
from apache_beam.ml.anomaly.base import AnomalyDetector
from apache_beam.ml.anomaly.base import AnomalyPrediction
from apache_beam.ml.anomaly.specifiable import specifiable
from apache_beam.ml.inference.base import KeyedModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import PredictionT

KeyT = TypeVar('KeyT')


@specifiable
class OfflineDetector(AnomalyDetector):
  """A offline anomaly detector that uses a provided model handler for scoring.

  Args:
    keyed_model_handler: The model handler to use for inference.
      Requires a `KeyModelHandler[Any, Row, PredictionT, Any]` instance.
    run_inference_args: Optional arguments to pass to RunInference
    **kwargs: Additional keyword arguments to pass to the base
      AnomalyDetector class.
  """
  @staticmethod
  def score_prediction_adapter(
      keyed_prediction: tuple[KeyT, PredictionResult]
  ) -> tuple[KeyT, AnomalyPrediction]:
    """Extracts a float score from `PredictionResult.inference` and wraps it.

    Takes a keyed `PredictionResult` from common ModelHandler output, assumes
    its `inference` attribute is a float-convertible score, and returns the key
    paired with an `AnomalyPrediction` containing that float score.

    Args:
      keyed_prediction: tuple of `(key, PredictionResult)`. `PredictionResult`
        must have an `inference` attribute supporting float conversion.

    Returns:
      tuple of `(key, AnomalyPrediction)` with the extracted score.

    Raises:
      AssertionError: If `PredictionResult.inference` doesn't support float().
    """

    key, prediction = keyed_prediction
    score = prediction.inference
    assert isinstance(score, SupportsFloat)
    return key, AnomalyPrediction(score=float(score))

  @staticmethod
  def label_prediction_adapter(
      keyed_prediction: tuple[KeyT, PredictionResult]
  ) -> tuple[KeyT, AnomalyPrediction]:
    """Extracts an integer label from `PredictionResult.inference` and wraps it.

    Takes a keyed `PredictionResult`, assumes its `inference` attribute is an
    integer-convertible label, and returns the key paired with an
    `AnomalyPrediction` containing that integer label.

    Args:
      keyed_prediction: tuple of `(key, PredictionResult)`. `PredictionResult`
        must have an `inference` attribute supporting int conversion.

    Returns:
      tuple of `(key, AnomalyPrediction)` with the extracted label.

    Raises:
      AssertionError: If `PredictionResult.inference` doesn't support int().
    """

    key, prediction = keyed_prediction
    label = prediction.inference
    assert isinstance(label, SupportsInt)
    return key, AnomalyPrediction(label=int(label))

  def __init__(
      self,
      keyed_model_handler: KeyedModelHandler[Any, beam.Row, PredictionT, Any],
      run_inference_args: Optional[dict[str, Any]] = None,
      **kwargs):
    super().__init__(**kwargs)

    # TODO: validate the model handler type
    self._keyed_model_handler = keyed_model_handler
    self._run_inference_args = run_inference_args or {}

    # always override model_identifier with model_id from the detector
    self._run_inference_args["model_identifier"] = self._model_id

  def learn_one(self, x: beam.Row) -> None:
    """Not implemented since OfflineDetector invokes RunInference directly."""
    raise NotImplementedError

  def score_one(self, x: beam.Row) -> Optional[float]:
    """Not implemented since OfflineDetector invokes RunInference directly."""
    raise NotImplementedError