# generated by datamodel-codegen:
#   filename:  inferences.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 arthur.client.rest.common.models import ResultRow
from arthur.client.rest.enrichments.models import ExplanationInput
from arthur.client.rest.models.models import Stage
from arthur.client.types import ByteField
from pydantic import BaseModel, Field, constr


class NewInference(BaseModel):
    """
    An inference to save
    """

    inference_timestamp: str
    """
    UTC timestamp of when the inference was made in ISO8601 string format
    """
    partner_inference_id: str
    """
    A unique id to associate with the inference. This is needed if you wish to later update the inference with ground truth and provides a quick way to retrieve the inference later.
    """
    batch_id: Optional[str] = None
    """
    Optional batch that the inference should be associated with
    """
    inference_data: Dict[str, Any]
    """
    The inference data to ingest. This should be a map of the attribute name to the attribute value. If ground truth data is included in the request, it should be added to the "ground_truth_data" field, not here.
    """
    ground_truth_timestamp: Optional[str] = None
    """
    Required if ground_truth_data is provided, this is the UTC timestamp to associate with the ground truth data
    """
    ground_truth_data: Optional[Dict[str, Any]] = None
    """
    Map of the ground truth attribute name to the attribute value
    """


class InferenceAttributeData(BaseModel):
    attribute_id: Optional[str] = None
    """
    The auto-generated unique UUID for the attribute associated with this datapoint.
    """
    attribute_name: Optional[constr(max_length=255)] = None
    """
    Name of the attribute associated with this datapoint.
    """
    data_value: Optional[Union[float, str, bool]] = Field(None, example=1.0)
    """
    value of the attribute for a specific inference.
    """
    stage: Optional[Stage] = None


class Type(str, Enum):
    """
    The type of dataset.
    """

    Batch = 'batch'
    Reference = 'reference'


class SageMakerDataFailures(BaseModel):
    """
    Failure reason for a single reference record
    """

    message: str = Field(..., example='missing field')
    """
    an error message
    """
    status: float = Field(..., example=400)
    """
    HTTP status code for the result
    """
    event_id: str = Field(..., example='unique-event-id')
    """
    The SageMaker event_id corresponding to the failed to ingest inference.
    """
    inference_index_in_event: int = Field(..., example=2)
    """
    The index of the inference within the SageMaker event, in case this event had multiple inferences under the same event_id.
    """


class SageMakerInferenceData(BaseModel):
    """
    File containing SageMaker Data Capture JSONL to bulk upload.
    """

    class Config:
        arbitrary_types_allowed = True

    inference_data: ByteField
    """
    A SageMaker Data Capture JSONL file containing inferences. The filename must end with \".jsonl\". Its data fields can either be encoded in either a CSV or a JSON format. If a user opts for a JSON encoding, the model inputs must be of format `{"instances": [{"features": [feature0, feature1, ...]}, {"features": [...]}]}` and the model outputs must be of format `{"predictions":[[prediction0, prediction1, ...], [prediction0, prediction1...]]}` for each JSON object in the JSON line file.
    """


class InferenceResult(BaseModel):
    """
    Result for a single inference
    """

    message: str = Field(..., example='success')
    """
    either "success" or an error message
    """
    status: float = Field(..., example=200)
    """
    HTTP status code for the result
    """


class FileInferenceData(BaseModel):
    """
    File containing inferences to bulk upload and optional batch_id OR file containing ground truth data. Either inferences.parquet or ground_truths.parquet must be included in the request.
    """

    class Config:
        arbitrary_types_allowed = True

    inference_data: Optional[ByteField] = None
    """
    Parquet or JSON file containing inferences. If JSON, follow the same format as the /inferences POST endpoint.
    """
    ground_truth_data: Optional[ByteField] = None
    """
    Parquet file containing ground truth, must contain the partner_inference_id field and ground_truth_timestamp field.
    """
    batch_id: Optional[str] = None
    """
    Optional batch_id to append to inferences in the provided file
    """
    image_data: Optional[ByteField] = None
    """
    For input type = Image models, this can be a single image file or a zip of multiple image files. inference_data is also required when this field is specified. The values of the image attribute in inference_data must match this file's filename for a single inference. When sending multiple inferences, the image attributes must match the paths of the corresponding images in the zipped image_data. The number of images must match the number of inferences in inference_data.
    """


class NewInferences(BaseModel):
    """
    New inferences to save
    """

    __root__: List[NewInference]
    """
    New inferences to save
    """


class PartnerIdInferenceGroundTruth(BaseModel):
    """
    The inference ground truth
    """

    ground_truth_timestamp: str
    """
    The UTC timestamp to associate with the ground truth data
    """
    partner_inference_id: str
    """
    The pre-defined unique id associated with the inference.
    """
    ground_truth_data: Dict[str, Any]
    """
    Map of the ground truth attribute name to the attribute value
    """


class Inference(BaseModel):
    id: constr(max_length=255)
    """
    The auto-generated unique UUID for the inference
    """
    partner_inference_id: Optional[constr(max_length=255)] = None
    """
    User defined unique inference id
    """
    timestamp: datetime
    """
    Inference UTC timestamp in ISO8601 format
    """
    anomaly_score: Optional[float] = None
    """
    This score shows anomalous an inference is and will fall between 0 and 1. A lower score suggests the inference is more anomalous/a larger outlier.
    """
    data: List[InferenceAttributeData]
    explanation: Optional[ExplanationInput] = None


class ImageInferenceEnum(str, Enum):
    RawImage = 'raw_image'
    ResizedImage = 'resized_image'
    Thumbnail = 'thumbnail'
    LIMEExplanation = 'lime_explanation'


class ReferenceDataRequest(BaseModel):
    """
    Reference set data to upload.
    """

    class Config:
        arbitrary_types_allowed = True

    reference_data_parquet: Optional[ByteField] = Field(
        None, alias='reference_data.parquet'
    )
    """
    Parquet file containing reference set data. The schema of the file must match the model attributes.
    """
    reference_data: Optional[ByteField] = None
    """
    Json or parquet file containing reference set data. The schema of the file must match the model attributes.
    """
    image_data: Optional[ByteField] = None
    """
    For input type = Image models, this can be a single image file or a zip of multiple image files. reference_data.parquet is also required when this field is specified. The values of the image attribute in reference_data.parquet must match this file's filename for a single inference. When sending multiple inferences, the image attributes must match the paths of the corresponding images in the zipped image_data. The number of images must match the number of rows in reference_data.parquet.
    """


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 ReferenceDataFailures(BaseModel):
    """
    Failure reason for a single reference record
    """

    message: str = Field(..., example='missing field')
    """
    an error message
    """
    status: float = Field(..., example=400)
    """
    HTTP status code for the result
    """


class ClosedRequest(BaseModel):
    status: str = Field(..., example='uploaded')
    """
    Setting the status to "uploaded" will close the batch and start the dataset post-processing.
    """
    total_record_count: Optional[int] = None
    """
    A total count of all the records in this dataset across all calls to ingestion service. This is intended to come from the customer so we can verify no data was lost in transmission.
    """


class SuccessResponse(BaseModel):
    message: Optional[str] = Field(None, example='success')


class DatasetStatusEnum(str, Enum):
    """
    The status of the dataset.
    """

    Started = 'started'
    WaitingForDataLoadToFinish = 'waiting_for_data_load_to_finish'
    Uploaded = 'uploaded'
    Processing = 'processing'
    Completed = 'completed'
    Failed = 'failed'


class QueryResult(BaseModel):
    query_result: List[ResultRow]


class BatchResponse(BaseModel):
    model_id: str
    """
    The unique UUID for the model.
    """
    batch_id: str
    """
    The id of the batch. Only returned for "batch" type datasets.
    """
    records: int
    """
    The number of records in the dataset
    """
    partner_record_count: int
    """
    The total number of records sent by the customer. This is an optional field that is populated by the "total_record_count" in the dataset POST request.
    """
    status: DatasetStatusEnum
    created_at: str
    """
    Timestamp the dataset was created in the Arthur platform.
    """
    updated_at: str
    """
    Timestamp the dataset was last updated in the Arthur platform.
    """


class DatasetResponse(BaseModel):
    model_id: str
    """
    The unique UUID for the model.
    """
    type: Type
    """
    The type of dataset.
    """
    batch_id: Optional[str] = None
    """
    The id of the batch. Only returned for "batch" type datasets.
    """
    records: int
    """
    The number of records in the dataset
    """
    partner_record_count: Optional[int] = None
    """
    The total number of records sent by the customer. This is an optional field that is populated by the "total_record_count" in the dataset POST request.
    """
    status: DatasetStatusEnum
    created_at: Optional[str] = None
    """
    Timestamp the dataset was created in the Arthur platform.
    """
    updated_at: Optional[str] = None
    """
    Timestamp the dataset was last updated in the Arthur platform.
    """


class InferencesResponse(BaseModel):
    """
    Statuses for each inference
    """

    counts: StatusCounts
    results: List[InferenceResult]
    """
    An array of result objects for each inference. The order matches the original request order
    """


class DatasetsResponse(BaseModel):
    __root__: List[DatasetResponse]


class SageMakerDataResponse(BaseModel):
    """
    Summary of failures
    """

    counts: StatusCounts
    failures: List[SageMakerDataFailures]
    """
    An array of result objects for each inference. The order matches the original request order
    """


class InferencesGroundTruth(BaseModel):
    """
    List of inference ground truth data
    """

    __root__: List[PartnerIdInferenceGroundTruth]
    """
    List of inference ground truth data
    """


class ReferenceDatasetResponse(BaseModel):
    model_id: Optional[str] = None
    """
    The unique UUID for the model.
    """
    records: Optional[int] = None
    """
    The number of records in the dataset
    """
    partner_record_count: Optional[int] = None
    """
    The total number of records sent by the customer. This is an optional field that is populated by the "total_record_count" in the dataset POST request.
    """
    status: Optional[DatasetStatusEnum] = None
    created_at: Optional[str] = None
    """
    Timestamp the dataset was created in the Arthur platform.
    """
    updated_at: Optional[str] = None
    """
    Timestamp the dataset was last updated in the Arthur platform.
    """


class ReferenceDataResponse(BaseModel):
    """
    Summary of failures
    """

    counts: StatusCounts
    failures: List[ReferenceDataFailures]
    """
    An array of result objects for each inference. The order matches the original request order
    """
