# Copyright 2021-2025 Kolena Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import field
from typing import Dict
from typing import List
from typing import Literal
from typing import Union

from kolena._utils.pydantic_v1 import StrictBool
from kolena._utils.pydantic_v1 import StrictStr
from kolena._utils.pydantic_v1.dataclasses import dataclass


@dataclass(frozen=True)
class CategoricalBucket:
    value: Union[StrictStr, StrictBool, int, None]
    count: int


@dataclass(frozen=True)
class NumericBucket:
    lower: float
    upper: float
    count: int
    interval_type: Literal["closed", "right-open"] = "right-open"


@dataclass(frozen=True)
class CategoricalStats:
    histogram: List[CategoricalBucket]
    data_type: Literal["categorical"] = "categorical"


@dataclass(frozen=True)
class NumericStats:
    count: int
    mean: float
    median: float
    min: float
    max: float
    stddev: float
    sum: float
    histogram: List[NumericBucket]
    data_type: Literal["numeric"] = "numeric"


@dataclass(frozen=True)
class ArraySizeStats:
    mean: float
    median: float
    min: int
    max: int
    stddev: float
    sum: int
    histogram: List[CategoricalBucket]


@dataclass(frozen=True)
class ArrayStats:
    length: ArraySizeStats
    data_type: Literal["array"] = "array"


@dataclass(frozen=True)
class FieldStats:
    key: str
    data: Union[CategoricalStats, NumericStats, ArrayStats]


@dataclass(frozen=True)
class SingleStatsResponse:
    datapoint: List[FieldStats] = field(default_factory=list)
    # key by model id first, then eval_config_id
    result: Dict[str, Dict[str, List[FieldStats]]] = field(default_factory=dict)
    llm: List[FieldStats] = field(default_factory=list)


@dataclass(frozen=True)
class FieldStatsResponse:
    stats: List[SingleStatsResponse]
