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

from __future__ import annotations

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

from arthur.client.rest.common.models import Page, PageSize, TotalCount, TotalPages
from arthur.client.rest.enrichments.models import ExplainabilityEnrichmentRequest
from pydantic import BaseModel, Field, conint, constr


class ModelHealthScore(BaseModel):
    model_id: str = Field(..., example='8519dce3-d1fd-4791-9f17-c15f43e43dcf')
    """
    id of the model
    """
    health_score: float = Field(..., example=78.9823)
    """
    the health score of the model between 0-100
    """
    timestamp: str = Field(..., example='2021-05-21 00:00:00+00:00')
    """
    the time the health score was calculated in UTC
    """


class ValueType(str, Enum):
    """
    Attribute value type.
    """

    Boolean = 'BOOLEAN'
    Float = 'FLOAT'
    Image = 'IMAGE'
    Integer = 'INTEGER'
    String = 'STRING'
    Timestamp = 'TIMESTAMP'
    BoundingBox = 'BOUNDING_BOX'
    UnstructuredText = 'UNSTRUCTURED_TEXT'
    Tokens = 'TOKENS'
    TokenLikelihoods = 'TOKEN_LIKELIHOODS'
    RankedList = 'RANKED_LIST'
    StringArray = 'ARRAY(STRING)'
    TimeSeries = 'TIME_SERIES'


class InputType(str, Enum):
    """
    The input type of the model.
    """

    Tabular = 'TABULAR'
    NLP = 'NLP'
    Image = 'IMAGE'
    TimeSeries = 'TIME_SERIES'


class OutputType(str, Enum):
    """
    The output type of the model.
    """

    Multiclass = 'MULTICLASS'
    Multilabel = 'MULTILABEL'
    Regression = 'REGRESSION'
    ObjectDetection = 'OBJECT_DETECTION'
    TokenSequence = 'TOKEN_SEQUENCE'
    RankedList = 'RANKED_LIST'


class ModelExpand(str, Enum):
    Attributes = 'attributes'
    Explainability = 'explainability'
    OrganizationID = 'organization_id'
    ModelGroupName = 'model_group_name'
    InferenceStats = 'inference_stats'
    CriticalAlertCount = 'critical_alert_count'
    HealthScores = 'health_scores'
    ExpectedThroughputGBPerDay = 'expected_throughput_gb_per_day'
    IsBinaryClasifier = 'is_binary_classifier'
    HasReferenceSetData = 'has_reference_set_data'
    HasGroundTruthInferences = 'has_ground_truth_inferences'


class Type(str, Enum):
    """
    What kind of entity does the ID represent.
    """

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


class CreatedBy(BaseModel):
    """
    Contains the information about an entity.
    """

    id: str
    """
    ID of the entity
    """
    type: Type
    """
    What kind of entity does the ID represent.
    """


class ModelHealthResponse(BaseModel):
    results: List[ModelHealthScore]


class AttributeBin(BaseModel):
    """
    A list of the attribute's bins. An attribute will only have bins if it is not categorical. The bin start is exclusive and the end is inclusive, (continuous_start, continuous_end]. Use Null to represent an open end of a bin.
    """

    continuous_start: Optional[float] = None
    """
    The exclusive start value for the bucket.
    """
    continuous_end: Optional[float] = None
    """
    The inclusive end value for the bucket.
    """


class AttributeCategory(BaseModel):
    """
    A list of the attribute's categories. An attribute will only have categories if it is marked as categorical.
    """

    label: Optional[str] = None
    """
    A label for the category.
    """
    value: str
    """
    The category value.
    """


class ExplainabilityResponse(BaseModel):
    """
    Describes the explainability configuration of a model. The fields marked as 'expandable' will only be returned when we see the query param 'expand=explainability'.
    """

    model_server_cpu: Optional[str] = None
    """
    Expandable. The CPU units upper bound to apply to the explainability server. If not set, the container can use all of the CPU available on the node. Maximum precision is 1m.
    """
    model_server_memory: Optional[str] = None
    """
    Expandable. The memory setting for the explainability server. If not set, will default to '1Gi'.
    """
    explanation_nsamples: Optional[int] = None
    """
    Expandable. Number of samples to use when generating the explanation. For SHAP, this corresponds to the 'nsamples' parameter, which defaults to len(features) + 2048. For LIME, this corresponds to 'num_samples', which defaults to 5000.
    """
    explanation_algo: Optional[str] = None
    """
    The algorithm to use to generate explanations. If explainability is enabled, this value will always be returned. We currently support 'lime' and 'shap' and default to 'lime'.
    """
    enabled: bool
    """
    Flag that indicates whether explainability is enabled for the model, defaults to false.
    """


class Stage(str, Enum):
    """
    Stage of this attribute in the model pipeline.
    """

    GroundTruth = 'GROUND_TRUTH'
    GroundTruthClass = 'GROUND_TRUTH_CLASS'
    NonInputData = 'NON_INPUT_DATA'
    PipelineInput = 'PIPELINE_INPUT'
    PredictedValue = 'PREDICTED_VALUE'
    PredictFunctionInput = 'PREDICT_FUNCTION_INPUT'


class Status(str, Enum):
    """
    The on-boarding status of the model
    """

    Pending = 'Pending'
    Creating = 'Creating'
    Ready = 'Ready'
    CreationFailed = 'CreationFailed'
    Archiving = 'Archiving'
    Archived = 'Archived'
    ArchiveFailed = 'ArchiveFailed'
    Unknown = 'Unknown'


class ExpandEnum(str, Enum):
    Attributes = 'attributes'
    Explainability = 'explainability'
    OrganizationID = 'organization_id'
    ExpectedThroughputGBPerDay = 'expected_throughput_gb_per_day'
    IsBinaryClassifier = 'is_binary_classifier'
    HasReferenceSetData = 'has_reference_set_data'
    HasGroundTruthInferences = 'has_ground_truth_inferences'


class IDTypeEnum(str, Enum):
    ID = 'id'
    PartnerModelID = 'partner_model_id'


class Tag(BaseModel):
    id: Optional[str] = None
    """
    UUID of the tag
    """
    organization_id: Optional[str] = None
    """
    UUI of the organization
    """
    name: str
    """
    Name of the tag
    """


class TagUpdateRequest(BaseModel):
    name: str
    """
    New name to update tag with
    """


class PaginatedTagResponse(BaseModel):
    """
    Paginated list of tags.
    """

    data: Optional[List[Tag]] = None
    page: Optional[Page] = None
    page_size: Optional[PageSize] = None
    total_pages: Optional[TotalPages] = None
    total_count: Optional[TotalCount] = None


class ModelAttribute(BaseModel):
    """
    A model attribute. Note that an attribute can only have categories if it is categorical and bins if it is not categorical, it can never have both.
    """

    id: Optional[str] = None
    """
    Auto-generated unique UUID for the attribute.
    """
    name: constr(max_length=255)
    """
    Name of the attribute.
    """
    label: Optional[constr(max_length=255)] = None
    """
    Label for attribute. If attribute has an encoded name, a more readable label can be set.
    """
    value_type: ValueType
    stage: Stage
    position: Optional[conint(ge=0)] = Field(None, example=0)
    """
    The array position of attribute within the stage. Required in the `PIPELINE_INPUT` stage.
    """
    categorical: bool
    """
    Flag that indicates whether or not this attribute is categorical.
    """
    min_range: Optional[Union[int, float]] = None
    """
    Optional field that indicates the minimum attribute value.
    """
    max_range: Optional[Union[int, float]] = None
    """
    Optional field that indicates the maximum attribute value.
    """
    monitor_for_bias: Optional[bool] = False
    """
    Indicates whether this attribute should be included in bias tracking.
    """
    categories: Optional[List[AttributeCategory]] = None
    """
    If the attribute is categorical, this will contain the attribute's categories.
    """
    bins: Optional[List[AttributeBin]] = None
    """
    Contains the attribute's bins. Only continuous attributes can have bins.
    """
    is_unique: Optional[bool] = False
    """
    Flag to set if this attribute should always be unique.
    """
    is_positive_predicted_attribute: Optional[bool] = False
    """
    Flag for binary classification models to specify positive predicted. attribute.
    """
    attribute_link: Optional[str] = None
    """
    For predicted value or ground truth staged attributes this should be populated with a corresponding predicted value or ground truth attribute name.
    """
    token_attribute_link: Optional[str] = None
    """
    For UNSTRUCTURED_TEXT, TOKENS, or TOKEN_LIKELIHOOD valued attributes, this should be populated with a corresponding attribute name of the tokenized text attribute or the raw text attribute.
    """
    gt_class_link: Optional[str] = None
    """
    For predicted value, optional mapping to corresponding ground truth class values.
    """
    implicit: Optional[bool] = False
    """
    True if this attribute was created implicitly by the API.
    """


class ModelRequest(BaseModel):
    display_name: Optional[constr(max_length=255)] = None
    """
    An optional display name for the model.
    """
    partner_model_id: constr(max_length=255)
    """
    Client provided unique id to associate with the model. This field must be unique across all active models cannot be changed once set!
    """
    input_type: InputType
    output_type: OutputType
    is_batch: Optional[bool] = Field(None, example=False)
    """
    Boolean value to determine whether the model sends inferences in batch or streaming format. Defaults to False.
    """
    description: Optional[constr(max_length=255)] = None
    """
    Optional description of the model.
    """
    model_group_id: Optional[str] = None
    """
    The UUID for the model group.
    """
    version_label: Optional[str] = None
    """
    Model version.
    """
    attributes: Optional[List[ModelAttribute]] = Field(None, min_items=1)
    """
    Attribute definitions for the model. Included in response only if the query param 'expand=attributes' is present.
    """
    pixel_height: Optional[int] = None
    """
    Image height for CV bounding box model images
    """
    pixel_width: Optional[int] = None
    """
    Image width for CV bounding box model images
    """
    image_class_labels: Optional[List[str]] = None
    """
    the categories a CV model can classify images or bounding boxes into
    """
    tags: Optional[List[str]] = None
    """
    Keywords to associate with the model.
    """
    classifier_threshold: Optional[float] = None
    """
    Threshold value for classification models, default is 0.5.
    """
    text_delimiter: Optional[str] = Field(None, example=',')
    """
    Only valid for models with `input_type` equal to `NLP`. Represents the text delimiter to divide input strings.
    """
    explainability_enabled: Optional[ExplainabilityEnrichmentRequest] = None
    accuracy_enabled: Optional[bool] = None
    """
    whether the accuracy metric should be enabled for the model
    """


class ModelObject(BaseModel):
    id: str
    """
    The auto-generated unique UUID for the model.
    """
    display_name: Optional[constr(max_length=255)] = None
    """
    An optional display name for the model.
    """
    partner_model_id: constr(max_length=255)
    """
    Client provided unique id to associate with the model. This field must be unique across all active models and cannot be changed once set!
    """
    description: Optional[constr(max_length=255)] = None
    """
    Optional description of the model.
    """
    input_type: InputType
    output_type: OutputType
    is_batch: Optional[bool] = Field(None, example=False)
    """
    Boolean value to determine whether the model sends inferences in batch or streaming format. Defaults to False.
    """
    archived: Optional[bool] = None
    """
    Indicates whether or not a model has been archived, defaults to false.
    """
    created_at: Optional[datetime] = None
    """
    UTC timestamp in ISO8601 format of when the model was created.
    """
    updated_at: Optional[datetime] = None
    """
    UTC timestamp in ISO8601 format of when the model was last updated.
    """
    attributes: Optional[List[ModelAttribute]] = Field(None, min_items=1)
    """
    Attribute definitions for the model. Included only if the query param 'expand=attributes' is present.
    """
    organization_id: Optional[str] = None
    """
    The UUID of the organization.
    """
    tags: Optional[List[str]] = None
    """
    Keywords to associate with the model.
    """
    classifier_threshold: Optional[float] = None
    """
    Threshold value for classification models, default is 0.5.
    """
    text_delimiter: Optional[str] = Field(None, example=',')
    """
    Only valid for models with `input_type` equal to `NLP`. Represents the text delimiter to divide input strings.
    """
    status: Optional[Status] = Field(None, example='Ready')
    """
    The on-boarding status of the model
    """
    model_group_id: Optional[str] = None
    """
    The UUID for the model group.
    """
    version_label: Optional[str] = None
    """
    Model version.
    """
    created_by: CreatedBy


class PaginatedAttributeResponse(BaseModel):
    data: List[ModelAttribute]
    """
    List of model objects.
    """
    page: Page
    page_size: PageSize
    total_pages: Optional[TotalPages] = None
    total_count: Optional[TotalCount] = None


class ModelResponse(ModelObject):
    explainability: Optional[ExplainabilityResponse] = None
    model_group_name: Optional[str] = None
    """
    The name of the model group the model belongs to.
    """
    total_inference_count: Optional[int] = None
    """
    The total number of inferences uploaded for the model.
    """
    average_inference_count: Optional[float] = None
    """
    The average of the daily inference counts for the model over the last month.
    """
    last_inference_timestamp: Optional[str] = None
    """
    The timestamp of the last inference for the model.
    """
    health_score: Optional[float] = Field(None, example=78.9823)
    """
    The aggregated health score of the model between 0-100. For details see Arthur Docs.
    """
    accuracy_health_score: Optional[float] = Field(None, example=86.5)
    """
    The health score of the model accuracy between 0-100. For details on how the score is calculated see Arthur Docs.
    """
    drift_health_score: Optional[float] = Field(None, example=90.076)
    """
    The health score of the model's data drift between 0-100. For details on how the score is calculated see Arthur Docs.
    """
    ingestion_health_score: Optional[float] = Field(None, example=76.976)
    """
    The health score of the model's data ingestion between 0-100. For details on how the score is calculated see Arthur Docs.
    """
    critical_alert_count: Optional[float] = None
    """
    The total number of critical alerts for the model that are active (status is new or acknowledged).
    """


class ModelUpdateRequest(BaseModel):
    display_name: Optional[constr(max_length=255)] = None
    """
    Display name of the model to update.
    """
    description: Optional[constr(max_length=255)] = None
    """
    Description of model to update.
    """
    attributes: Optional[List[ModelAttribute]] = None
    """
    Attribute definitions to update on the model. Only ground_truth attributes can be updated after inferences have been sent. Before inferences have been sent any attributes can be updated. The attributes provided in this class will replace the ones currently associated with the model.
    """
    tags: Optional[List[str]] = None
    """
    Tags to set on the model object.
    """
    overwrite_tags: Optional[bool] = None
    """
    If true, will overwrite the existing tags for this model with the tags specified in this request. Defaults to false.
    """
    classifier_threshold: Optional[float] = None
    """
    Threshold value for classification models, default is 0.5. This can only be updated before any inferences have been sent
    """


class PaginatedModelResponse(BaseModel):
    data: List[ModelResponse]
    """
    List of model objects.
    """
    page: Page
    page_size: PageSize
    total_pages: Optional[TotalPages] = None
    total_count: Optional[TotalCount] = None
