from typing import Optional, Union, List

from pydantic import Field

from anssformats.formatbasemodel import FormatBaseModel
from anssformats.source import Source as SourceFormat


class Prediction(FormatBaseModel):
    """A generic prediction from an analytical model (AI/ML, statistical, rule-based, etc.)

    This class provides a flexible structure for representing any type of prediction
    with associated metrics and model provenance.

    Attributes
    ----------
    label: string identifying what is being predicted (e.g., "phase", "magnitude",
        "eventType", "custom_attribute")

    value: the predicted value - can be string, number, or structured object

    probability: optional float [0.0-1.0] containing the probability of this prediction

    metrics: optional dict containing additional prediction metrics such as confidence,
        uncertainty quantification (std, ranges, credible intervals), entropy, etc.
        Structure is flexible to support any model-specific metrics.

    modelID: optional string identifying which model made this prediction

    modelVersion: optional string containing the version of the model

    source: optional Source object containing the source/author of the
    """

    label: str = Field(
        ...,
        description="What is being predicted (e.g., phase, magnitude, eventType, distance)",
    )
    value: Union[str, float, int, list, dict] = Field(
        ..., description="The predicted value"
    )
    probability: Optional[float] = Field(
        None, ge=0.0, le=1.0, description="Probability of this prediction [0.0-1.0]"
    )
    metrics: Optional[dict] = Field(
        None,
        description="Additional prediction metrics (confidence, uncertainty, entropy, etc.)",
    )
    modelID: Optional[str] = Field(
        None, description="Identifier for the model that made this prediction"
    )
    modelVersion: Optional[str] = Field(None, description="Version of the model")
    source: Optional[SourceFormat] = Field(
        None, description="Source/author of the model"
    )


class Analytics(FormatBaseModel):
    """A conversion class used to create, parse, and validate analytical information
    from models such as pickers, analytical models, AI/ML, etc.

    This class provides an extensible structure supporting multiple models, arbitrary
    prediction labels, and custom extensions.

    Attributes
    ----------

    predictions: optional list of Prediction objects containing predictions from one
        or more models. Each prediction can represent any type of analytical output
        with associated confidence metrics and model provenance.

    extensions: optional dict containing custom key-value pairs for experimental code and debugging.

    """

    predictions: Optional[List[Prediction]] = Field(
        None, description="Array of predictions from one or more analytical models"
    )

    extensions: Optional[dict] = Field(
        None, description="Custom key-value pairs for experimental or debuging data"
    )
