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

from __future__ import annotations

from enum import Enum
from typing import List, Optional

from arthur.client.rest.common.models import (
    DataDriftRequestBase,
    QueryFilters,
    QueryGroupBy,
    ResultRow,
)
from pydantic import BaseModel


class QueryResult(BaseModel):
    query_result: List[ResultRow]
    sampling_threshold: Optional[float] = None


class DistributionsAttribute(BaseModel):
    property: str
    """
    The name of the property to select. This can be the name of any attribute associated with the model.
    """
    alias: Optional[str] = None
    """
    If an alias is provided, this will be returned as the name for the associated result.
    """
    num_bins: Optional[float] = None
    """
    The number of "buckets" to distribute the attribute into. Value must be at least 2 and no greater than 200. If not specified, we will default to 20.
    """


class DistributionValue(BaseModel):
    function: Optional[str] = None
    """
    optional, an aggregation function like count, sum, abs, max, min, avg
    """
    property: str
    """
    The name of the property to select. This can be the name of any attribute associated with the model.
    """
    alias: Optional[str] = None
    """
    If an alias is provided, this will be returned as the name for the associated result.
    """


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

    minute = 'minute'
    hour = 'hour'
    day = 'day'
    month = 'month'
    year = 'year'
    batch_id = 'batch_id'


class DataDriftTableRequest(BaseModel):
    property: str
    """
    This will be the name of one of the model's attributes or an enriched attribute. Required if alias is not provided.
    """
    num_bins: Optional[int] = None
    """
    Continuous attributes only - Sets the number of bins in which to perform the (PSI)[https://scholarworks.wmich.edu/cgi/viewcontent.cgi?article=4249&context=dissertations] 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
    """
    base: DataDriftRequestBase
    target: DataDriftRequestBase
    group_by: Optional[QueryGroupBy] = None
    rollup: Optional[Rollup] = None
    """
    An optional parameter that will aggregate the calculated PSI value by the supported dimension.
    """


class DistributionsRequest(BaseModel):
    x: DistributionsAttribute
    y: DistributionsAttribute
    values: Optional[List[DistributionValue]] = None
    filter: Optional[QueryFilters] = None


class Rollup1(str, Enum):
    minute = 'minute'
    hour = 'hour'
    day = 'day'
    month = 'month'
    year = 'year'
    batch_id = 'batch_id'


class PSIBucketValue(BaseModel):
    bucket: Optional[str] = None
    rollup: Optional[Rollup1] = None
    group_by_property: Optional[str] = None
    base_bucket_max: Optional[float] = None
    base_bucket_min: Optional[float] = None
    base_count_per_bucket: Optional[float] = None
    base_ln_probability_per_bucket: Optional[float] = None
    base_probability_per_bucket: Optional[float] = None
    base_total: Optional[float] = None
    target_bucket_max: Optional[float] = None
    target_bucket_min: Optional[float] = None
    target_count_per_bucket: Optional[float] = None
    target_ln_probability_per_bucket: Optional[float] = None
    target_probability_per_bucket: Optional[float] = None
    target_total: Optional[float] = None
    probability_difference: Optional[float] = None
    ln_probability_difference: Optional[float] = None
    psi: Optional[float] = None


class PsiValue(BaseModel):
    rollup: Optional[Rollup1] = None
    group_by_property: Optional[str] = None
    psi_value: Optional[float] = None


class DataDriftResponse(BaseModel):
    query_result: Optional[List[PsiValue]] = None


class DataDriftTableResponse(BaseModel):
    query_result: Optional[List[PSIBucketValue]] = None
