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

from __future__ import annotations

from enum import Enum
from typing import Dict, List, Optional, Union

from pydantic import BaseModel, Extra, Field


class Rollup(str, Enum):
    """
    An optional parameter that will aggregate the calculated data drift value by the supported dimension.
    """

    Minute = 'minute'
    Hour = 'hour'
    Day = 'day'
    Month = 'month'
    Year = 'year'
    BatchID = 'batch_id'


class Source(str, Enum):
    """
    Location of the data (inference or reference)
    """

    Inference = 'inference'
    Reference = 'reference'


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

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


class Page(BaseModel):
    __root__: int = Field(..., example=1)


class PageSize(BaseModel):
    __root__: int = Field(..., example=100)


class ResultRow(BaseModel):
    """
    Represents 1 row of result data
    """

    pass

    class Config:
        extra = Extra.allow


class Comparator(str, Enum):
    """
    gt = greater than, gte = greater than or equal, lt = less than, lte = less than or equal, eq = equal, like = used for text search
    """

    GreaterThan = 'gt'
    GreaterThanOrEqual = 'gte'
    LessThan = 'lt'
    LessThanOrEqual = 'lte'
    Equal = 'eq'
    In = 'in'
    Like = 'like'


class QueryFilter(BaseModel):
    alias: Optional[str] = None
    """
    This will map back to a defined function in the select field and allows you to filter on a calculated function. Required if property is not provided.
    """
    comparator: Comparator
    """
    gt = greater than, gte = greater than or equal, lt = less than, lte = less than or equal, eq = equal, like = used for text search
    """
    property: Optional[str] = None
    """
    The name of the property to filter on. This field is required if alias is not provided.
    """
    value: Union[int, float, str]
    """
    The value to compare the property against
    """
    value_type: Optional[str] = None


class QueryFilters(BaseModel):
    __root__: List[QueryFilter]


class QueryGroupByItem(BaseModel):
    alias: Optional[str] = None
    """
    This should map back to an alias defined in the select clause. Required if property is not provided.
    """
    property: Optional[str] = None
    """
    This will be the name of one of the model's attributes or an enriched attribute. Required if alias is not provided.
    """


class QueryOrderByItem(BaseModel):
    alias: Optional[str] = None
    """
    This should map back to an alias defined in the select clause. Required if property is not provided.
    """
    direction: Optional[str] = None
    """
    Either 'asc' or 'desc'. Defaults to 'asc'.
    """
    property: Optional[str] = None
    """
    This will be the name of one of the model's attributes or an enriched attribute. Required if alias is not provided.
    """


class From(str, Enum):
    """
    Data to query. Inferences are the raw inferences sent to the platform. Enriched includes the raw inferences, plus enriched properties like explanations and anomaly scores. Reference is the reference data set uploaded to the platform. The default is 'inference'.
    """

    Inference = 'inference'
    Enriched = 'enriched'
    Reference = 'reference'


class Sort(BaseModel):
    __root__: str
    """
    Must be supplied in the format [column_name] to denote asc sort by this column OR -[column_name] to denote desc sort by this column
    """


class StatusCounts(BaseModel):
    """
    Counts of the results by status
    """

    failure: float
    """
    The number of failed inferences in the request
    """
    success: float
    """
    The number of successful inferences in the request
    """
    total: float
    """
    The total number of inferences in the request
    """


class TotalCount(BaseModel):
    __root__: int = Field(..., example=2)
    """
    The total number of results the query returned
    """


class TotalPages(BaseModel):
    __root__: int = Field(..., example=1)
    """
    The number of result pages that this query generated, the ceiling of the total number of results divided by the page size
    """


class From1(str, Enum):
    """
    Data to query. Inferences are the raw inferences sent to the platform. Enriched includes the raw inferences, plus enriched properties like explanations and anomaly scores. Reference is the reference data set uploaded to the platform. The default is 'inference'.
    """

    inference = 'inference'
    enriched = 'enriched'
    reference = 'reference'


class AnyArray(BaseModel):
    __root__: List[Union[int, float, str, KeyValueObject]]


class DataDriftRequestBase(BaseModel):
    filter: Optional[QueryFilters] = None
    source: Optional[Source] = None
    """
    Location of the data (inference or reference)
    """


class QueryGroupBy(BaseModel):
    __root__: List[QueryGroupByItem]


class QueryOrderBy(BaseModel):
    __root__: List[QueryOrderByItem]


class DataDriftRequest(BaseModel):
    base: DataDriftRequestBase
    group_by: Optional[QueryGroupBy] = None
    metric: str
    """
    A required parameter that will determine type of data drift value.
    """
    num_bins: Optional[int] = None
    """
    Continuous attributes only - Sets the number of bins in which to perform the data drift calculation.  Bin boundaries are created by assigning each data point of the continuous property to 1 of [numBins] quantiles. For example: If you have an attribute of 100 values from 1 to 100, and passed numBins = 10, then 10 bins would be created with boundaries at 10, 20, 30, ..., 100
    """
    properties: List[str]
    """
    This is a list of model attributes or an enriched attributes. At least one is required.
    """
    rollup: Optional[Rollup] = None
    """
    An optional parameter that will aggregate the calculated data drift value by the supported dimension.
    """
    target: DataDriftRequestBase
    predicted_classes: Optional[List[str]] = None
    ground_truth_classes: Optional[List[str]] = None
    sample: Optional[float] = None
    disable_auto_sample: Optional[bool] = None


class NestedParameterReference(BaseModel):
    """
    Can be used in a function parameter to refer to a column by alias or nest another function as input. Only one of nested_function or alias_ref can be specified at a time.
    """

    alias_ref: Optional[str] = None
    """
    references another alias column by name
    """
    nested_function: Optional[QuerySelectItem] = None


class QueryRequest(BaseModel):
    filter: Optional[QueryFilters] = None
    from_: Optional[From] = Field(None, alias='from')
    """
    Data to query. Inferences are the raw inferences sent to the platform. Enriched includes the raw inferences, plus enriched properties like explanations and anomaly scores. Reference is the reference data set uploaded to the platform. The default is 'inference'.
    """
    group_by: Optional[QueryGroupBy] = None
    order_by: Optional[QueryOrderBy] = None
    select: QuerySelect
    subquery: Optional[QueryRequest] = None
    limit: Optional[int] = None
    offset: Optional[int] = None
    sample: Optional[float] = None
    disable_auto_sample: Optional[bool] = None


class QuerySelect(BaseModel):
    __root__: List[QuerySelectItem]


class QuerySelectItem(BaseModel):
    """
    Describes the fields and functions the query should return
    """

    alias: Optional[str] = None
    """
    If an alias is provided, this will be returned as the name for the associated result. By default, functions will have the alias set to the function name unless one is provided.
    """
    function: Optional[str] = None
    """
    The name of the function to use. Either property or function can be specified but not both. For a list of functions see https://docs.arthur.ai/docs/querying-performance
    """
    parameters: Optional[SelectParameter] = None
    property: Optional[str] = None
    """
    The name of the property to select. This can be the name of any attribute associated with the model, any of the enriched attributes ['anomaly_score', 'lime_feature_importance', 'shap_feature_importance'], or "*", which selects all of the model's non-enriched attributes. Either property or function can be specified but not both.
    """


class SelectParameter(BaseModel):
    """
    Some functions will require additional parameters, for example, you will want to specify the "threshold" for the count function for a classifier. See the Query Documentation for more details
    """

    class Config:
        extra = Extra.allow

    __root__: Optional[
        Dict[
            str,
            Union[float, str, bool, AnyArray, KeyValueObject, NestedParameterReference],
        ]
    ] = None


class MetricQuery(BaseModel):
    """
    This should be in the format described in our [Query Guide](/api-documentation/query_guide.html). Your query must have exactly one select and may not use the `group_by` or `subquery` fields. Don't include any batch or timestamp filters as these will be added dynamically. For details [see the full metrics docs](https://docs.arthur.ai/user-guide/walkthroughs/metrics_alerts.html#creating-custom-metrics).
    """

    select: Optional[QuerySelect] = None
    from_: Optional[From1] = Field(None, alias='from')
    """
    Data to query. Inferences are the raw inferences sent to the platform. Enriched includes the raw inferences, plus enriched properties like explanations and anomaly scores. Reference is the reference data set uploaded to the platform. The default is 'inference'.
    """
    filter: Optional[QueryFilters] = None
    group_by: Optional[QueryGroupBy] = None
    order_by: Optional[QueryOrderBy] = None


NestedParameterReference.update_forward_refs()
QueryRequest.update_forward_refs()
QuerySelect.update_forward_refs()
QuerySelectItem.update_forward_refs()
