# generated by datamodel-codegen:
#   filename:  metrics.yaml
#   timestamp: 2023-12-15T18:09:27+00:00

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Union
from uuid import UUID

from arthur.client.rest.common.models import (
    MetricQuery,
    Page,
    PageSize,
    QueryFilter,
    TotalCount,
    TotalPages,
)
from pydantic import BaseModel, Extra, Field


class CategoricalValueLimits(BaseModel):
    """
    Possible values within fixed set
    """

    categories: List
    """
    The list of allowed values. Not required if it can be inferred.
    """


class MetricGroupBy(BaseModel):
    __root__: str = Field(..., example='Education')
    """
    Model attribute to group by. If provided, the results will be broken down by each value or bin in the attribute
    """


class Value(BaseModel):
    class Config:
        extra = Extra.allow

    rollup: Optional[str] = Field(None, example='2022-08-03T00:00:00+00:00')
    """
    The time aggregation associated with this metric value
    """


class Rollup(str, Enum):
    """
    What period/dataset to aggregate over.
    """

    Minute = 'MINUTE'
    Hour = 'HOUR'
    Day = 'DAY'
    Month = 'MONTH'
    Year = 'YEAR'
    BatchID = 'BATCH_ID'
    Reference = 'REFERENCE'


class Type(str, Enum):
    """
    Type associated with the metric to create. If this field is not supplied the metric type will automatically be filled in.
    """

    ModelOutputMetric = 'model_output_metric'
    ModelPerformanceMetric = 'model_performance_metric'
    ModelDataDriftMetric = 'model_data_drift_metric'
    ModelDataBoundMetric = 'model_data_bound_metric'


class BaseMetric(BaseModel):
    """
    A base metric object used by both MetricRequest and MetricResponse
    """

    name: str = Field(..., example='Accuracy Rate')
    """
    The name for your custom metric, this should not match any other metric name for the model and cannot match any Arthur default metric names.
    """
    endpoint: str = Field(
        ...,
        example='/api/v3/models/4d722041-7274-6875-7220-53616d75656c/inferences/query',
    )
    """
    The relative endpoint to hit when executing this metric
    """
    query: MetricQuery
    type: Optional[Type] = None
    """
    Type associated with the metric to create. If this field is not supplied the metric type will automatically be filled in.
    """


class LastUpdatedByType(str, Enum):
    """
    Type of user that created or last updated this metric. Null if the metric was auto-generated by Arthur.
    """

    ServiceAccount = 'service-account'
    ArthurManaged = 'arthur-managed'
    IDPManaged = 'idp-managed'


class RangeValueLimits(BaseModel):
    """
    Possible values between a minimum and maximum value
    """

    min: float = Field(..., example=0.0)
    """
    The minimum value allowed. Not required if it can be inferred.
    """
    max: float = Field(..., example=1.0)
    """
    The maximum value allowed. Not required if it can be inferred.
    """


class MetricParameterValues(BaseModel):
    class Config:
        extra = Extra.allow

    __root__: Optional[Dict[str, Union[List[Union[float, str]], float, str]]] = None


class Bin(BaseModel):
    """
    the bin start and end, if group is continuous
    """

    continuous_start: Optional[float] = None
    """
    the start of the bin
    """
    continuous_end: Optional[float] = None
    """
    the end of the bin
    """


class GroupLegend(BaseModel):
    """
    Specification of group for a metric result.
    """

    name: Optional[str] = None
    """
    the name of the group, if group is categorical
    """
    bin: Optional[Bin] = None
    """
    the bin start and end, if group is continuous
    """


class ValueLimit(BaseModel):
    range: Optional[RangeValueLimits] = None


class ValueLimit1(BaseModel):
    categorical: Optional[CategoricalValueLimits] = None


class ParameterType(str, Enum):
    """
    The type of a supplied parameter
    """

    Attribute = 'ATTRIBUTE'
    Integer = 'INTEGER'
    Float = 'FLOAT'
    String = 'STRING'
    Timestamp = 'TIMESTAMP'


class ExpandMetricEnumEnum(str, Enum):
    Type = 'type'


class ExpandMetricEnum(BaseModel):
    __root__: List[ExpandMetricEnumEnum]


class MetricTypeEnum(str, Enum):
    ModelOutputMetric = 'model_output_metric'
    ModelDataBoundMetric = 'model_input_data_metric'
    ModelPerformanceMetric = 'model_performance_metric'
    ModelDataDriftMetric = 'model_data_drift_metric'


class MetricEvaluationRequest(BaseModel):
    """
    The arguments to use when evaluating a metric
    """

    rollup: Optional[Rollup] = None
    group_by: Optional[MetricGroupBy] = None
    parameters: Optional[MetricParameterValues] = None
    filters: Optional[List[QueryFilter]] = None


class MetricResult(BaseModel):
    """
    The result of a metric evaluation on a set of parameters, filters, and group
    """

    name: str = Field(..., example='accuracy_rate')
    """
    The name or alias of the function computed in the underlying query. A value with this key will be present in each item/map in the 'values' array in the response.
    """
    values: List[Value] = Field(
        ...,
        example=[
            {'rollup': '2022-08-03T00:00:00+00:00', 'accuracy_rate': 0.82},
            {'rollup': '2022-08-04T00:00:00+00:00', 'accuracy_rate': 0.76},
            {'rollup': '2022-08-05T00:00:00+00:00', 'accuracy_rate': 0.8},
        ],
    )
    """
    the value of the metric, aggregating by the rollup
    """
    parameters: Optional[Dict[str, Union[float, str]]] = Field(
        None, example={'classifier_threshold': 0.5}
    )
    """
    the parameters used when evaluating the metric
    """
    group: Optional[Dict[str, GroupLegend]] = Field(
        None, example={'Education': {'name': 'HS'}}
    )
    """
    the grouping applied
    """


class MetricEvaluationResponse(BaseModel):
    """
    The result of a metric evaluation
    """

    name: str = Field(..., example='Accuracy Rate')
    """
    The name of the metric
    """
    filters: Optional[List[QueryFilter]] = None
    """
    The dataset filters used during metric evaluation
    """
    results: List[MetricResult]
    """
    The results of metric evaluation
    """


class MetricParameterRequest(BaseModel):
    """
    A Parameter specification for a Metric
    """

    name: str = Field(..., example='classifier_threshold')
    parameter_type: ParameterType
    value_limits: Union[ValueLimit, ValueLimit1]
    """
    Acceptable values for the parameter. Exactly one of `categorical` or `range` must be specified.
    """
    default: Any = Field(..., example=0.5)
    """
    Default value(s) for the parameter. If not provided, it will be inferred if possible.
    """


class MetricRequest(BaseMetric):
    """
    A Metric containing a template query and parameter definitions
    """

    parameters: Optional[List[MetricParameterRequest]] = None


class MetricParameterResponse(MetricParameterRequest):
    id: Optional[UUID] = None


class MetricResponse(BaseMetric):
    id: Optional[UUID] = None
    valid_rollups: Optional[List[str]] = Field(
        None, example=['MINUTE', 'HOUR', 'DAY', 'MONTH', 'REFERENCE']
    )
    """
    The allowed "rollup" parameters for this metric
    """
    parameters: Optional[List[MetricParameterResponse]] = None
    last_updated_by: Optional[str] = None
    """
    User ID of the user that created or last updated this metric. Null if the metric was auto-generated by Arthur.
    """
    last_updated_by_type: Optional[LastUpdatedByType] = None
    """
    Type of user that created or last updated this metric. Null if the metric was auto-generated by Arthur.
    """
    created_at: Optional[datetime] = None
    """
    UTC timestamp of when the metric was created
    """
    updated_at: Optional[datetime] = None
    """
    UTC timestamp of when the metric was last updated
    """


class PaginatedModelMetricsResponse(BaseModel):
    """
    Paginated list of metrics.
    """

    metrics: List[MetricResponse]
    page: Page
    page_size: PageSize
    total_pages: Optional[TotalPages] = None
    total_count: Optional[TotalCount] = None
