import json
import math
import re
import threading
from abc import ABC, abstractmethod
from collections import defaultdict
from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass, field, fields, replace
from datetime import datetime, timedelta, timezone
from enum import Enum, IntEnum, auto, unique
from functools import cached_property
from itertools import chain, groupby, repeat, starmap
from random import random
from typing import (
    Any,
    BinaryIO,
    Generic,
    NamedTuple,
    Optional,
    TypeVar,
    Union,
    cast,
    get_args,
    overload,
)
from uuid import uuid4
from weakref import ProxyType, proxy

import numpy as np
import numpy.typing as npt
import pandas as pd
from pandas.core.dtypes.common import (
    is_bool_dtype,
    is_datetime64_any_dtype,
    is_datetime64tz_dtype,
    is_numeric_dtype,
)
from typing_extensions import TypeAlias, TypeGuard
from wrapt import ObjectProxy

from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX
from phoenix.datetime_utils import floor_to_minute


class DimensionRole(IntEnum): ...


@unique
class InvalidRole(DimensionRole):
    UNASSIGNED = auto()


@unique
class SingularDimensionalRole(DimensionRole):
    """Roles that cannot be assigned to more than one dimension."""

    # The (integer) ordering here is important in that it'll be used
    # as tie-breaker when e.g. the user assigns a column to both prediction
    # label and predicton score, in which case the role with a lower
    # integer value will prevail.
    PREDICTION_ID = 1 + max(InvalidRole)
    TIMESTAMP = auto()
    PREDICTION_LABEL = auto()
    PREDICTION_SCORE = auto()
    ACTUAL_LABEL = auto()
    ACTUAL_SCORE = auto()
    # Large Language Model (LLM) prompt and response pairs
    PROMPT = auto()
    RESPONSE = auto()


@unique
class MultiDimensionalRole(DimensionRole):
    # It's important to keep the numeric values disjoint among all subclass
    # enums (hence the +1 here), because we'll use `groupby(sorted(...))` to
    # collate the dimensions. The (integer) ordering here is also important
    # in that it'll be used as tie-breaker when e.g. the user assigns a
    # column to both feature and tag, in which case the role with a lower
    # integer value will prevail.
    FEATURE = 1 + max(SingularDimensionalRole)
    TAG = auto()


# Global variables for ease of making references.
PREDICTION_ID = SingularDimensionalRole.PREDICTION_ID
TIMESTAMP = SingularDimensionalRole.TIMESTAMP

PREDICTION_LABEL = SingularDimensionalRole.PREDICTION_LABEL
PREDICTION_SCORE = SingularDimensionalRole.PREDICTION_SCORE
ACTUAL_LABEL = SingularDimensionalRole.ACTUAL_LABEL
ACTUAL_SCORE = SingularDimensionalRole.ACTUAL_SCORE

PROMPT = SingularDimensionalRole.PROMPT
RESPONSE = SingularDimensionalRole.RESPONSE

FEATURE = MultiDimensionalRole.FEATURE
TAG = MultiDimensionalRole.TAG


def _rand_str() -> str:
    """Generates a random string, useful for adding a new column to a
    dataframe, such that it won't conflict with the existing column names.
    It's intended to be short and gibberish looking.
    """
    return hex(int(random() * 1e9))


@dataclass(frozen=True)
class SchemaSpec(ABC):
    """Superclass for the Schema itself and any of its components, all of
    which share the trait that they use column names to define the model
    structure (or substructures)."""

    def __post_init__(self) -> None:
        """Ensure column names are string at run time, which may not be true
        if the user takes them directly from `pd.dataframe.columns`, which
        can contain numbers. Phoenix always uses string to refer to columns
        and dimensions, so this step is intended to eliminate any potential
        issues caused by any pd.DataFrame having numbers as column names.
        """
        for f in fields(self):
            if f.type is str or set(get_args(f.type)) == {str, type(None)}:
                # Ensure string if type is `str` or `Optional[str]` (the
                # latter being `Union[str, None]` under the hood). Use
                # `__setattr__` because `self` is read-only, i.e. frozen.
                object.__setattr__(self, f.name, _coerce_str(getattr(self, f.name)))


@dataclass(frozen=True)
class CompositeDimensionSpec(SchemaSpec, ABC):
    """A dimension referencing multiple columns (besides its primary column).
    E.x. Embedding can point to 3 columns
    """

    @abstractmethod
    def __str__(self) -> str:
        """This lets the spec behave like a string in some situations,
        because it's basically a column name, but with extra stuff.
        """
        ...

    @abstractmethod
    def __iter__(self) -> Iterator[str]: ...


@dataclass(frozen=True)
class Embedding(CompositeDimensionSpec):
    vector: str
    """The column name of the vector values."""
    raw_data: Optional[str] = None
    link_to_data: Optional[str] = None
    display_name: Optional[str] = None

    def __str__(self) -> str:
        return self.vector

    def __iter__(self) -> Iterator[str]:
        for value in (self.vector, self.raw_data, self.link_to_data):
            if isinstance(value, str) and value:
                yield value


@dataclass(frozen=True)
class RetrievalEmbedding(Embedding):
    context_retrieval_ids: Optional[str] = None
    context_retrieval_scores: Optional[str] = None

    def __iter__(self) -> Iterator[str]:
        yield from super().__iter__()
        for value in (
            self.context_retrieval_ids,
            self.context_retrieval_scores,
        ):
            if isinstance(value, str) and value:
                yield value


class InferencesRole(Enum):
    """A dataframe's role in a Model: primary or reference (as
    baseline for drift).
    """

    PRIMARY = auto()
    REFERENCE = auto()


PRIMARY = InferencesRole.PRIMARY
REFERENCE = InferencesRole.REFERENCE


@dataclass(frozen=True, repr=False, eq=False)
class _ConstantValueSeriesFactory:
    """The intent is to share memory for readonly situations where a constant
    value, e.g. NaN, is expected from a non-existent dataframe column by the
    downstream calculations. To avoid duplicating memory when we require such
    columns multiple times or need to subset them frequently, we first check for
    a previously allocated array that's longer, and subset from it, thereby
    sharing the memory already allocated (since the value is constant and does
    not change). If the cached array is not long enough, we then allocate and
    cache a new array, replacing the old one. This cache is intended to be
    attached to a Model instance, and will be garbage collected alongside it,
    while the `lru_cache` decorator's cache is scoped at the module level and
    can potentially leak memory.
    """

    value: Any = field(default=np.nan)
    _dtype: Any = field(init=False, default=None)
    _cached_array: npt.NDArray[Any] = field(
        init=False,
        default_factory=lambda: np.empty(0),
    )
    """If a longer Series is requested, the cached array is expanded;
    otherwise, a subset can be returned, assuming it won't be altered by the
    caller.
    """
    _lock: threading.Lock = field(
        init=False,
        default_factory=threading.Lock,
    )
    """A lock is applied at the class instance level for thread safety, with
    minimal overhead expected, unless too many callers simultaneously rely on
    the same instance.
    """

    def __post_init__(self) -> None:
        if isinstance(self.value, float) and math.isnan(self.value):
            object.__setattr__(self, "_dtype", np.float32)

    def __call__(self, length: int) -> "pd.Series[Any]":
        with self._lock:
            if length > len(self._cached_array):
                object.__setattr__(
                    self,
                    "_cached_array",
                    np.full(length, self.value, dtype=self._dtype),
                )
            return pd.Series(self._cached_array[:length])


_Key = TypeVar("_Key", bound=Hashable)
_Value = TypeVar("_Value")


@dataclass(frozen=True, repr=False, eq=False)
class _Cache(Generic[_Key, _Value]):
    """A thread-safe type-safe generic cache backed by a dictionary.

    Example
    -------
    >>> c = _Cache[str, int]()
    >>> with c() as cache:
    >>>     cache["1"] = 2
    >>> with c() as cache:
    >>>     print(cache["1"])
    2
    """

    _cache: dict[_Key, _Value] = field(
        init=False,
        default_factory=dict,
    )
    _lock: threading.Lock = field(
        init=False,
        default_factory=threading.Lock,
    )

    @contextmanager
    def __call__(self) -> Iterator[dict[_Key, _Value]]:
        with self._lock:
            yield self._cache


DataFrameOrSeries: TypeAlias = Union[pd.DataFrame, "pd.Series[Any]"]
"""Either a table or a single row of data. Using a series to represent one
row of data (instead of a one-row dataframe) ensures that only scalar values
are extracted (instead of getting a column back)."""


@dataclass(frozen=True)
class Column:
    """Extracts a value (i.e. scalar) from pd.Series or a series (i.e. a
    column) from pd.DataFrame. If not found, return the default value
    (e.g. NaN) or a series of the default value (on each row), respectively.
    """

    name: str = ""
    is_dummy: bool = False
    """dummy columns are fillers for the model structure, so our functions
    can remain relatively declarative and compact, i.e. by not having to keep
    writing `if ... is not None:` everywhere. Dummy columns always return NaNs,
    because they have random column names not found in any dataframe."""
    _default: _ConstantValueSeriesFactory = field(
        default_factory=_ConstantValueSeriesFactory,
    )

    def __post_init__(self) -> None:
        if not self.name:
            object.__setattr__(self, "is_dummy", True)
            object.__setattr__(self, "name", _rand_str())

    @overload
    def __call__(self, data: pd.DataFrame) -> "pd.Series[Any]": ...

    @overload
    def __call__(self, data: "pd.Series[Any]") -> Any: ...

    def __call__(self, data: DataFrameOrSeries) -> Any:
        """Extracts a value from series, or a series from a dataframe. If
        not found, return the default value (e.g. NaN) or a series of the
        default value (on each row), respectively.
        """
        if isinstance(data, pd.DataFrame):
            try:
                return data.loc[:, self.name]
            except KeyError:
                # It's important to glue the index to the default series,
                # so it would look like the series came from the dataframe.
                return self._default(len(data)).set_axis(data.index)
        if isinstance(data, pd.Series):
            try:
                return data.at[self.name]
            except KeyError:
                return self._default.value
        raise ValueError(f"invalid data: {repr(data)}")

    def __iter__(self) -> Iterator[str]:
        """This is to partake in the iteration of column names by a
        larger data structure of which this object is a member. Dummy
        columns need not be yielded because they represent columns
        that don't (physically) exist (i.e. they are just fillers for
        the Model to make the Schema whole).
        """
        if not self.is_dummy:
            yield self.name

    def __str__(self) -> str:
        return self.name


class DataType(IntEnum):
    UNKNOWN = auto()
    DISCRETE = auto()  # usually names, e.g. California, or Monday.
    CONTINUOUS = auto()  # usually amounts, e.g. 5 bucks, or 6.7 miles.


UNKNOWN = DataType.UNKNOWN
DISCRETE = DataType.DISCRETE
CONTINUOUS = DataType.CONTINUOUS


@dataclass(frozen=True)
class Dimension(Column, ABC):
    """All Dimensions are Columns (i.e. dataframe column) but not all Columns
    are Dimensions, e.g. some are just ancillary, like a column of URLs.
    """

    role: DimensionRole = InvalidRole.UNASSIGNED
    """The dimension's role in the Model, e.g. FEATURE, TIME, or PREDICTION_ID.
    Must not be unassigned at initialization."""
    data_type: DataType = UNKNOWN
    """Denotes whether the dimension values are continuous or discrete, useful
    for binning."""
    _model: Optional["ProxyType[Model]"] = field(repr=False, default=None)
    """Holds a weak reference to the parent Model, if any. This is the channel
    for talking to actual data, and the Model can help out if needed."""

    def __post_init__(self) -> None:
        super().__post_init__()
        if isinstance(self.role, InvalidRole):
            # This is to let the superclass be initialized without arguments.
            # But we really want the role to be specified for a Dimension.
            raise ValueError("role must be assigned")

    def __getitem__(self, df_role: InferencesRole) -> "pd.Series[Any]":
        if self._model is None:
            return pd.Series(dtype=object)
        model = cast(Model, self._model)
        return self(model[df_role])

    @property
    def display_name(self) -> str:
        """In general, a dimension's name is that of its primary column. But
        dimensions with special roles, e.g. `PREDICTION_ID` have names
        that match their roles, e.g. `Prediction ID` (without underscores),
        and don't necessarily match the names of the actual underlying columns
        in the dataframes, which e.g. could be a derived column with a random
        name such as in the case of the TIME dimension, which refers to a
        normalized version of the actual timestamps column (if it existed).
        """
        if isinstance(self.role, SingularDimensionalRole):
            return _title_case_no_underscore(self.role.name)
        return self.name


@dataclass(frozen=True)
class ScalarDimension(Dimension):
    @property
    def min_max(self) -> tuple[Any, Any]:
        if self._model is None:
            return np.nan, np.nan
        model = cast(Model, self._model)
        return model.dimension_min_max_from_all_df(self.name)

    @property
    def categories(self) -> tuple[str, ...]:
        if self._model is None or self.data_type is CONTINUOUS:
            return ()
        model = cast(Model, self._model)
        return model.dimension_categories_from_all_inferences(self.name)


@dataclass(frozen=True)
class EmbeddingDimension(Dimension):
    link_to_data: Column = field(default_factory=Column)
    raw_data: Column = field(default_factory=Column)
    display_name: str = ""

    def __post_init__(self) -> None:
        super().__post_init__()
        if not self.display_name:
            object.__setattr__(self, "display_name", self.name)
        # optimization: share the default series factory
        object.__setattr__(
            self,
            "link_to_data",
            replace(self.link_to_data, _default=self._default),
        )
        # optimization: share the default series factory
        object.__setattr__(
            self,
            "raw_data",
            replace(self.raw_data, _default=self._default),
        )

    @classmethod
    def from_embedding(cls, emb: Embedding, **kwargs: Any) -> "EmbeddingDimension":
        # Use `from_embedding` instead of `__init__` because the latter is
        # needed by replace() and we don't want to clobber the generated
        # version.
        return cls(
            _coerce_str(emb.vector),
            link_to_data=Column(_coerce_str(emb.link_to_data)),
            raw_data=Column(_coerce_str(emb.raw_data)),
            display_name=_coerce_str(emb.display_name),
            **kwargs,
        )

    def __iter__(self) -> Iterator[str]:
        """This is to partake in the iteration of column names by a
        larger data structure of which this object is a member.
        """
        yield from super().__iter__()
        if not self.raw_data.is_dummy:
            yield from self.raw_data
        if not self.link_to_data.is_dummy:
            yield from self.link_to_data


@dataclass(frozen=True)
class RetrievalEmbeddingDimension(EmbeddingDimension):
    context_retrieval_ids: Column = field(default_factory=Column)
    context_retrieval_scores: Column = field(default_factory=Column)

    @classmethod
    def from_prompt_embedding(
        cls, emb: RetrievalEmbedding, **kwargs: Any
    ) -> "RetrievalEmbeddingDimension":
        # Use `from_embedding` instead of `__init__` because the latter is
        # needed by replace() and we don't want to clobber the generated
        # version.
        return cls(
            _coerce_str(emb.vector),
            link_to_data=Column(_coerce_str(emb.link_to_data)),
            raw_data=Column(_coerce_str(emb.raw_data)),
            display_name=_coerce_str(emb.display_name),
            context_retrieval_ids=Column(_coerce_str(emb.context_retrieval_ids)),
            context_retrieval_scores=Column(_coerce_str(emb.context_retrieval_scores)),
            **kwargs,
        )

    def __iter__(self) -> Iterator[str]:
        """This is to partake in the iteration of column names by a
        larger data structure of which this object is a member.
        """
        yield from super().__iter__()
        if not self.context_retrieval_ids.is_dummy:
            yield from self.context_retrieval_ids
        if not self.context_retrieval_scores.is_dummy:
            yield from self.context_retrieval_scores


Name: TypeAlias = str
ColumnKey: TypeAlias = Union[Name, Column, SingularDimensionalRole]
MultiDimensionKey: TypeAlias = Union[MultiDimensionalRole, Sequence[DimensionRole]]
RowNumbering: TypeAlias = Union[int, list[int]]


def _is_column_key(key: Any) -> TypeGuard[ColumnKey]:
    return isinstance(key, (str, Column, SingularDimensionalRole))


def _is_multi_dimension_key(
    key: Any,
) -> TypeGuard[MultiDimensionKey]:
    if isinstance(key, str):
        return False
    try:
        for k in iter(key):
            if not isinstance(k, DimensionRole):
                return False
        return True
    except TypeError:
        return isinstance(key, MultiDimensionalRole)


def _is_dimension_type_filter(
    key: Any,
) -> TypeGuard[tuple[MultiDimensionKey, Union[type[ScalarDimension], type[EmbeddingDimension]]]]:
    return (
        isinstance(key, tuple)
        and len(key) == 2
        and _is_multi_dimension_key(key[0])
        and isinstance(key[1], (Dimension, EmbeddingDimension))
    )


def _is_named_df(obj: Any) -> TypeGuard[tuple[Name, pd.DataFrame]]:
    return (
        isinstance(obj, tuple)
        and len(obj) == 2
        and isinstance(obj[0], str)
        and isinstance(obj[1], pd.DataFrame)
    )


RowId: TypeAlias = int


class TimeRange(NamedTuple):
    start: datetime
    stop: datetime


class ModelData(ObjectProxy, ABC):  # type: ignore
    """pd.DataFrame or pd.Series wrapped with extra functions and metadata."""

    def __init__(
        self,
        data: DataFrameOrSeries,
        /,
        _model: Optional["ProxyType[Model]"] = None,
    ) -> None:
        super().__init__(data)
        self._self_model = _model

    @property
    @abstractmethod
    def null_value(self) -> Any: ...

    def __getitem__(self, key: Any) -> Any:
        if _is_column_key(key):
            if isinstance(key, Column):
                return key(self)
            if self._self_model is None or self.empty:
                return self.null_value
            model = cast(Model, self._self_model)
            return model[key](self)
        return super().__getitem__(key)


class EventId(NamedTuple):
    """Identifies an event."""

    row_id: int = 0
    inferences_id: InferencesRole = PRIMARY

    def __str__(self) -> str:
        return ":".join(map(str, self))


class Event(ModelData):
    """pd.Series wrapped with extra functions and metadata."""

    def __init__(
        self,
        series: "pd.Series[Any]",
        /,
        event_id: EventId,
        **kwargs: Any,
    ) -> None:
        super().__init__(series, **kwargs)
        self._self_id = event_id

    @property
    def id(self) -> EventId:
        return self._self_id

    def null_value(self) -> float:
        return np.nan

    @overload
    def __getitem__(self, key: ColumnKey) -> Any: ...

    @overload
    def __getitem__(self, key: Any) -> Any: ...

    def __getitem__(self, key: Any) -> Any:
        return super().__getitem__(key)


class Events(ModelData):
    """pd.DataFrame wrapped with extra functions and metadata."""

    def __init__(
        self,
        df: pd.DataFrame,
        /,
        role: InferencesRole,
        **kwargs: Any,
    ) -> None:
        super().__init__(df, **kwargs)
        self._self_role = role

    @property
    def null_value(self) -> "pd.Series[Any]":
        return pd.Series(dtype=object)

    @cached_property
    def time_range(self) -> TimeRange:
        if self._self_model is None or self.empty:
            now = datetime.now(timezone.utc)
            return TimeRange(now, now)
        model = cast(Model, self._self_model)
        min_max = _agg_min_max(model[TIMESTAMP](self))
        start_time = cast(datetime, min_max.min())
        end_time = cast(datetime, min_max.max())
        # Add one minute to end_time, because time intervals are right
        # open and one minute is the smallest interval allowed.
        stop_time = end_time + timedelta(minutes=1)
        # Round down to the nearest minute.
        start = floor_to_minute(start_time)
        stop = floor_to_minute(stop_time)
        return TimeRange(start, stop)

    def __iter__(self) -> Iterator[Event]:
        for i, event in self.iterrows():
            yield Event(
                event,
                event_id=EventId(i, self._self_role),
                _model=self._self_model,
            )

    @overload
    def __getitem__(self, key: ColumnKey) -> "pd.Series[Any]": ...

    @overload
    def __getitem__(self, key: list[RowId]) -> "Events": ...

    def __getitem__(self, key: Any) -> Any:
        if isinstance(key, list):
            return Events(
                self.loc[key],
                role=self._self_role,
                _model=self._self_model,
            )
        return super().__getitem__(key)


class Inferences(Events):
    """pd.DataFrame wrapped with extra functions and metadata."""

    def __init__(
        self,
        df: pd.DataFrame,
        /,
        name: str,
        **kwargs: Any,
    ) -> None:
        super().__init__(df, **kwargs)
        self._self_name = name

    @property
    def name(self) -> str:
        return self._self_name

    @property
    def display_name(self) -> str:
        """
        Return a human-readable name for this dataset. If the user doesn't
        provide a name, the name is generated but this name is not human
        friendly. Falls back to the role of the dataset if no name is provided.
        """
        ds_name = self._self_name
        if ds_name.startswith(GENERATED_INFERENCES_NAME_PREFIX):
            # The generated names are UUIDs so use the role as the name
            return "primary" if self.role is InferencesRole.PRIMARY else "reference"
        return ds_name

    @property
    def role(self) -> InferencesRole:
        return self._self_role

    @property
    def is_empty(self) -> bool:
        return len(self) == 0

    @cached_property
    def primary_key(self) -> "pd.Index[Any]":
        return cast("pd.Index[Any]", pd.Index(self[PREDICTION_ID]))

    @overload
    def __getitem__(self, key: ColumnKey) -> "pd.Series[Any]": ...

    @overload
    def __getitem__(self, key: list[RowId]) -> Events: ...

    def __getitem__(self, key: Any) -> Any:
        if isinstance(key, list):
            return Events(
                self.iloc[key].set_axis(key),
                role=self._self_role,
                _model=self._self_model,
            )
        return super().__getitem__(key)


@dataclass(frozen=True, repr=False)
class Model:
    """A Model consists of a set of dataframes and a set of Dimensions.
    Dimensions are objects that extracts columns from dataframes. Not all
    dataframe columns are Dimensions. Some are ancillary, e.g. a column of
    URLs. All Dimensions have a primary column and its name is the same as
    the Dimension's. Some Dimensions may have more than one column, but only
    one is primary. When a Dimension tries to extract a non-existent column
    from a dataframe, a default (or fallback) column will be returned, e.g.
    a column of NaNs.
    """

    _inference_sets: dict[InferencesRole, Inferences]
    _dimensions: dict[Name, Dimension]
    _dim_names_by_role: dict[DimensionRole, list[Name]]
    _original_columns_by_role: dict[InferencesRole, "pd.Index[Any]"]
    _default_timestamps_factory: _ConstantValueSeriesFactory
    _nan_series_factory: _ConstantValueSeriesFactory
    _dimension_categories_from_all_inferences: _Cache[Name, tuple[str, ...]]
    _dimension_min_max_from_all_inferences: _Cache[Name, tuple[float, float]]

    def __init__(
        self,
        dimensions: Iterable[Dimension],
        dataframes: Iterable[Union[pd.DataFrame, tuple[Name, pd.DataFrame]]],
        /,
        treat_omitted_columns_as_features: bool = True,
        timestamps_already_normalized: bool = False,
        df_already_sorted_by_time: bool = False,
        # TODO: Consider moving validations here.
        df_already_validated: bool = False,
    ):
        # memoization
        object.__setattr__(
            self,
            "_dimension_categories_from_all_inferences",
            _Cache[Name, "pd.Series[Any]"](),
        )
        object.__setattr__(
            self,
            "_dimension_min_max_from_all_inferences",
            _Cache[Name, tuple[float, float]](),
        )

        df_names, dfs = cast(
            tuple[Iterable[Name], Iterable[pd.DataFrame]],
            zip(*_coerce_tuple(dataframes)),
        )
        str_col_dfs = _coerce_str_column_names(dfs)
        padded_dfs = _add_padding(str_col_dfs, pd.DataFrame)
        padded_df_names = _add_padding(df_names, _rand_str)
        inference_sets = starmap(
            self._new_inferences,
            zip(padded_dfs, padded_df_names, InferencesRole),
        )
        # Store inferences by role.
        object.__setattr__(
            self,
            "_inference_sets",
            {inferences.role: inferences for inferences in inference_sets},
        )
        # Preserve originals, useful for exporting.
        object.__setattr__(
            self,
            "_original_columns_by_role",
            {role: inferences.columns for role, inferences in self._inference_sets.items()},
        )

        object.__setattr__(
            self,
            "_nan_series_factory",
            _ConstantValueSeriesFactory(np.nan),
        )
        object.__setattr__(
            self,
            "_default_timestamps_factory",
            _ConstantValueSeriesFactory(datetime.now(timezone.utc)),
        )

        # Store dimensions by name. In general, a dimension's name is that of
        # its primary column. But dimensions with special roles, e.g.
        # `PREDICTION_ID` have names that match their roles, e.g.
        # `Prediction ID` (without underscores), and don't necessarily match
        # the names of actual underlying columns in the dataframes.
        object.__setattr__(
            self,
            "_dimensions",
            {dim.name: self._new_dimension(dim) for dim in dimensions},
        )
        if treat_omitted_columns_as_features:
            self._dimensions.update(
                (name, self._new_dimension(name, role=FEATURE))
                for name in _get_omitted_column_names(
                    self._dimensions.values(),
                    self._inference_sets.values(),
                )
            )

        # Group dimension names by roles. Store in a `defaultdict(list)`.
        object.__setattr__(
            self,
            "_dim_names_by_role",
            defaultdict(list, _group_names_by_dim_role(self._dimensions.values())),
        )

        # Guess the data type, i.e. continuous or discrete, for scalar
        # features and tags.
        for dim in self[ScalarDimension]:
            if dim.data_type is not UNKNOWN:
                continue
            self._dimensions[dim.name] = replace(
                dim,
                data_type=(
                    _guess_data_type(
                        dataset.loc[:, dim.name]
                        for dataset in self._inference_sets.values()
                        if dim.name in dataset.columns
                    )
                ),
            )

        # Add PREDICTION_ID if missing.
        # Add TIMESTAMP if missing.
        # If needed, normalize the timestamps values.
        # If needed, sort the dataframes by time.
        for inferences_role, dataset in list(self._inference_sets.items()):
            df = dataset.__wrapped__
            df_original_columns = self._original_columns_by_role[inferences_role]

            # PREDICTION_ID
            dim_pred_id = self._dimensions.get(
                next(iter(self._dim_names_by_role[PREDICTION_ID]), ""),
                self._new_dimension(PREDICTION_ID),
            )
            if dim_pred_id.name not in df_original_columns:
                df[dim_pred_id.name] = _series_uuid(len(df))
            self._dimensions[dim_pred_id.name] = dim_pred_id
            self._dim_names_by_role[PREDICTION_ID] = [dim_pred_id.name]

            # TIMESTAMP
            dim_time = self._dimensions.get(
                next(iter(self._dim_names_by_role[TIMESTAMP]), ""),
                self._new_dimension(TIMESTAMP),
            )
            if dim_time.name not in df_original_columns:
                df[dim_time.name] = self._default_timestamps_factory(len(df))
            else:
                if not timestamps_already_normalized:
                    # Don't clobber the original (or any other column).
                    # Store as new column with a random name.
                    new_name, old_name = _rand_str(), dim_time.name
                    df[new_name] = _normalize_timestamps(df[old_name])
                    dim_time = replace(dim_time, name=new_name)
                    del self._dimensions[old_name]
                if not df_already_sorted_by_time:
                    # Sort data for faster search.
                    df = df.sort_values(dim_time.name)
            self._dimensions[dim_time.name] = dim_time
            self._dim_names_by_role[TIMESTAMP] = [dim_time.name]
            # Set time column as index for use by pd.Grouper.
            df = df.set_index(dim_time.name, drop=False)

            # Update dataset since its dataframe may have changed.
            self._inference_sets[inferences_role] = self._new_inferences(
                df, name=dataset.name, role=inferences_role
            )

    @cached_property
    def is_empty(self) -> bool:
        """Returns True if the model has no data."""
        return not any(map(len, self._inference_sets.values()))

    def export_rows_as_parquet_file(
        self,
        row_numbers: Mapping[InferencesRole, Iterable[int]],
        parquet_file: BinaryIO,
        cluster_ids: Optional[Mapping[InferencesRole, Mapping[int, str]]] = None,
    ) -> None:
        """
        Given row numbers, exports dataframe subset into parquet file.
        Duplicate rows are removed. If the model hase more than one dataset, a
        new column is added to the dataframe containing the dataset name of
        each row in the exported data. The name of the added column will be
        `__phoenix_dataset_name__`.

        Parameters
        ----------
        row_numbers: Mapping[InferencesRole, Iterable[int]]
            mapping of dataset role to list of row numbers
        parquet_file: file handle
            output parquet file handle
        cluster_ids: Optional[Mapping[InferencesRole, Mapping[int, str]]]
            mapping of inferences role to mapping of row number to cluster id.
            If cluster_ids is non-empty, a new column is inserted to the
            dataframe containing the cluster IDs of each row in the exported
            data. The name of the added column name is `__phoenix_cluster_id__`.
        """
        export_dataframes = [pd.DataFrame()]
        model_has_multiple_inference_sets = (
            sum(not df.empty for df in self._inference_sets.values()) > 1
        )
        for inferences_role, numbers in row_numbers.items():
            df = self._inference_sets[inferences_role]
            columns = [
                df.columns.get_loc(column_name)
                for column_name in self._original_columns_by_role[inferences_role]
            ]
            rows = pd.Series(sorted(set(numbers)))
            filtered_df = df.iloc[rows, columns].reset_index(drop=True)
            if model_has_multiple_inference_sets:
                filtered_df["__phoenix_dataset_name__"] = df.display_name
            if cluster_ids and (ids := cluster_ids.get(inferences_role)):
                filtered_df["__phoenix_cluster_id__"] = rows.apply(ids.get)
            export_dataframes.append(filtered_df)
        pd.concat(export_dataframes).to_parquet(
            parquet_file,
            index=False,
            allow_truncated_timestamps=True,
            coerce_timestamps="ms",
        )

    @cached_property
    def scalar_dimensions(self) -> tuple[ScalarDimension, ...]:
        """Put these in a cached sequence because currently their positions
        in this list also determine their node IDs in graphql.
        """
        return tuple(
            dim
            for dim in self._dimensions.values()
            if not dim.is_dummy
            and dim.role not in (PREDICTION_ID, TIMESTAMP, RESPONSE)
            and isinstance(dim, ScalarDimension)
        )

    @cached_property
    def embedding_dimensions(self) -> tuple[EmbeddingDimension, ...]:
        """Put these in a cached sequence because currently their positions
        in this list also determine their node IDs in graphql.
        """
        return tuple(
            dim
            for dim in self._dimensions.values()
            if not dim.is_dummy and isinstance(dim, EmbeddingDimension)
        )

    def dimension_categories_from_all_inferences(
        self,
        dimension_name: Name,
    ) -> tuple[str, ...]:
        dim = self[dimension_name]
        if dim.data_type is CONTINUOUS:
            return cast(tuple[str, ...], ())
        with self._dimension_categories_from_all_inferences() as cache:
            try:
                return cache[dimension_name]
            except KeyError:
                pass
        categories_by_dataset = (
            pd.Series(dim[role].unique()).dropna().astype(str) for role in InferencesRole
        )
        all_values_combined = chain.from_iterable(categories_by_dataset)
        ans = tuple(np.sort(pd.Series(all_values_combined).unique()))
        with self._dimension_categories_from_all_inferences() as cache:
            cache[dimension_name] = ans
        return ans

    def dimension_min_max_from_all_df(
        self,
        dimension_name: Name,
    ) -> tuple[float, float]:
        dim = self[dimension_name]
        if dim.data_type is not CONTINUOUS:
            return (np.nan, np.nan)
        with self._dimension_min_max_from_all_inferences() as cache:
            try:
                return cache[dimension_name]
            except KeyError:
                pass
        min_max_by_df = (_agg_min_max(dim[df_role]) for df_role in InferencesRole)
        all_values_combined = chain.from_iterable(min_max_by_df)
        min_max = _agg_min_max(pd.Series(all_values_combined))
        ans = (min_max.min(), min_max.max())
        with self._dimension_min_max_from_all_inferences() as cache:
            cache[dimension_name] = ans
        return ans

    @overload
    def __getitem__(self, key: type[Inferences]) -> Iterator[Inferences]: ...

    @overload
    def __getitem__(self, key: InferencesRole) -> Inferences: ...

    @overload
    def __getitem__(self, key: ColumnKey) -> Dimension: ...

    @overload
    def __getitem__(self, key: MultiDimensionKey) -> Iterator[Dimension]: ...

    @overload
    def __getitem__(self, key: type[ScalarDimension]) -> Iterator[ScalarDimension]: ...

    @overload
    def __getitem__(self, key: type[EmbeddingDimension]) -> Iterator[EmbeddingDimension]: ...

    @overload
    def __getitem__(self, key: type[Dimension]) -> Iterator[Dimension]: ...

    @overload
    def __getitem__(
        self,
        key: tuple[
            MultiDimensionKey,
            Union[type[ScalarDimension], type[EmbeddingDimension]],
        ],
    ) -> Iterator[Dimension]: ...

    def __getitem__(self, key: Any) -> Any:
        if key is Inferences:
            return self._inference_sets.values()
        if isinstance(key, InferencesRole):
            return self._inference_sets[key]
        if _is_column_key(key):
            return self._get_dim(key)
        if _is_multi_dimension_key(key):
            return self._get_multi_dims(key)
        if _is_dimension_type_filter(key):
            return self._get_multi_dims_by_type(key)
        if key is ScalarDimension:
            return filter(
                lambda dim: type(dim) is ScalarDimension,
                self._dimensions.values(),
            )
        if key is EmbeddingDimension:
            return filter(
                lambda dim: type(dim) is EmbeddingDimension,
                self._dimensions.values(),
            )
        if key is Dimension:
            return self._dimensions.values()
        raise KeyError(f"invalid key: {repr(key)}")

    def _get_dim(self, key: ColumnKey) -> Dimension:
        if isinstance(key, DimensionRole):
            key = self._dim_names_by_role[key][0]
        if isinstance(key, str):
            return self._dimensions[key]
        raise KeyError(f"invalid key: {repr(key)}")

    def _get_multi_dims(self, key: MultiDimensionKey) -> Iterator[Dimension]:
        for k in _iterate_except_str(key):
            if isinstance(k, DimensionRole):
                for name in self._dim_names_by_role[k]:
                    yield self._dimensions[name]
            elif isinstance(k, str):
                yield self._dimensions[k]
            else:
                raise KeyError(f"invalid key: {repr(key)}")

    def _get_multi_dims_by_type(
        self,
        key: tuple[
            MultiDimensionKey,
            Union[type[ScalarDimension], type[EmbeddingDimension]],
        ],
    ) -> Iterator[Dimension]:
        return filter(lambda dim: type(dim) is key[1], self[key[0]])

    @overload
    def _new_dimension(
        self,
        obj: DimensionRole,
        cls: type[Dimension] = ScalarDimension,
        **kwargs: Any,
    ) -> Dimension: ...

    @overload
    def _new_dimension(
        self,
        obj: Name,
        cls: type[Dimension] = ScalarDimension,
        **kwargs: Any,
    ) -> Dimension: ...

    @overload
    def _new_dimension(
        self,
        obj: Dimension,
        **kwargs: Any,
    ) -> Dimension: ...

    def _new_dimension(
        self, obj: Any, cls: type[Dimension] = ScalarDimension, **kwargs: Any
    ) -> Dimension:
        """Creates a new Dimension or copies an existing one, setting the
        model weak reference to the `self` Model instance, and sharing the
        NaN series factory as an optimization.
        """
        if isinstance(obj, Name):
            return cls(
                obj,
                **kwargs,
                _model=proxy(self),
                _default=self._nan_series_factory,
            )
        if isinstance(obj, DimensionRole):
            return cls(
                role=obj,
                **kwargs,
                _model=proxy(self),
                _default=self._nan_series_factory,
            )
        if isinstance(obj, Dimension):
            return replace(
                obj,
                **kwargs,
                _model=proxy(self),
                _default=self._nan_series_factory,
            )
        raise ValueError(f"invalid argument: {repr(obj)}")

    def _new_inferences(
        self,
        df: pd.DataFrame,
        /,
        name: str,
        role: InferencesRole,
    ) -> Inferences:
        """Creates a new Inferences, setting the model weak reference to the
        `self` Model instance.
        """
        return Inferences(df, name=name, role=role, _model=proxy(self))


@dataclass(frozen=True)
class Schema(SchemaSpec):
    prediction_id: Optional[str] = None
    timestamp: Optional[str] = None
    prediction_label: Optional[str] = None
    prediction_score: Optional[str] = None
    actual_label: Optional[str] = None
    actual_score: Optional[str] = None
    prompt: Optional[Embedding] = None
    response: Optional[Union[str, Embedding]] = None
    features: Iterable[Union[str, CompositeDimensionSpec]] = field(default_factory=list)
    tags: Iterable[Union[str, CompositeDimensionSpec]] = field(default_factory=list)

    # internal attribute not exposed to users
    _dimensions: list[Dimension] = field(
        init=False, repr=False, hash=False, compare=False, default_factory=list
    )

    def __post_init__(self) -> None:
        super().__post_init__()
        object.__setattr__(self, "features", tuple(_iterate_except_str(self.features)))
        object.__setattr__(self, "tags", tuple(_iterate_except_str(self.tags)))
        grouped_by_name = groupby(
            sorted(self._make_dims(), key=lambda dim: (dim.name, dim.role)),
            key=lambda dim: dim.name,
        )
        for name, group in grouped_by_name:
            self._dimensions.append(next(group))

    def _make_dims(self) -> Iterator[Dimension]:
        """Iterate over all dimensions defined by the Schema, substituting
        with dummy dimensions for ones omitted by user. The dummy dimensions
        have randomly generated names that can change for each iteration, but
        currently there's no need to iterate more than once."""
        for spec, dimension_role, data_type in chain(
            (
                (self.prediction_id, PREDICTION_ID, DISCRETE),
                (self.timestamp, TIMESTAMP, CONTINUOUS),
                (self.prediction_label, PREDICTION_LABEL, DISCRETE),
                (self.prediction_score, PREDICTION_SCORE, CONTINUOUS),
                (self.actual_label, ACTUAL_LABEL, DISCRETE),
                (self.actual_score, ACTUAL_SCORE, CONTINUOUS),
                (self.prompt, PROMPT, DISCRETE),
                (self.response, RESPONSE, DISCRETE),
            ),
            zip(self.features, repeat(FEATURE), repeat(UNKNOWN)),
            zip(self.tags, repeat(TAG), repeat(UNKNOWN)),
        ):
            if not isinstance(spec, CompositeDimensionSpec):
                spec = _coerce_str(spec)
            assert isinstance(dimension_role, DimensionRole)  # for mypy
            if isinstance(spec, str):
                if dimension_role is PROMPT:
                    yield RetrievalEmbeddingDimension(
                        spec,
                        role=dimension_role,
                        data_type=data_type,
                    )
                else:
                    yield ScalarDimension(
                        spec,
                        role=dimension_role,
                        data_type=data_type,
                    )
            elif isinstance(spec, RetrievalEmbedding):
                yield RetrievalEmbeddingDimension.from_prompt_embedding(
                    spec,
                    role=dimension_role,
                    data_type=data_type,
                )
            elif isinstance(spec, Embedding):
                yield EmbeddingDimension.from_embedding(
                    spec,
                    role=dimension_role,
                    data_type=data_type,
                )
            else:
                raise TypeError(f"{dimension_role} has unrecognized type: {type(spec)}")

    def __call__(
        self,
        *dataframes: Union[pd.DataFrame, tuple[Name, pd.DataFrame]],
        **kwargs: Any,
    ) -> Model:
        """Dimensions are the "baton" that Schema hands over to Model."""
        _raise_if_too_many_dataframes(len(dataframes))
        return Model(iter(self._dimensions), dataframes, **kwargs)

    def __iter__(self) -> Iterator[str]:
        for f in fields(self):
            if not f.init:
                continue
            value = getattr(self, f.name)
            if isinstance(value, str) and value:
                yield value
                continue
            try:
                for v in iter(value):
                    if isinstance(v, str) and v:
                        yield v
                        continue
                    try:
                        yield from iter(v)
                    except TypeError:
                        pass
            except TypeError:
                pass

    def to_json(self) -> str:
        return json.dumps(_jsonify(self))

    @classmethod
    def from_json(cls, input: str) -> "Schema":
        data = json.loads(input)
        try:
            return cls(**{k: _objectify(v) for k, v in data.items()})
        except ValueError as e:
            raise ValueError(f"invalid json data: {repr(data)}") from e


def _agg_min_max(series: "pd.Series[Any]") -> "pd.Series[Any]":
    return series.agg(["min", "max"])


def _get_omitted_column_names(
    dimensions: Iterable[Dimension],
    dataframes: Iterable[pd.DataFrame],
) -> Iterator[str]:
    dataframe_columns = chain.from_iterable(df.columns for df in dataframes)
    schema_columns = chain.from_iterable(dimensions)
    yield from set(dataframe_columns) - set(schema_columns)


def _group_names_by_dim_role(
    dimensions: Iterable[Dimension],
) -> Iterator[tuple[DimensionRole, list[str]]]:
    return (
        (role, [dim.name for dim in dims])
        for role, dims in groupby(
            sorted(dimensions, key=lambda dim: dim.role),
            key=lambda dim: dim.role,
        )
    )


def _guess_data_type(series: Iterable["pd.Series[Any]"]) -> DataType:
    for s in series:
        if s.empty:
            continue
        if is_bool_dtype(s):
            break
        if is_numeric_dtype(s) or is_datetime64_any_dtype(s):
            return CONTINUOUS
    return DISCRETE


def _iterate_except_str(obj: Any) -> Iterator[Any]:
    """Strings are iterable (by character), but we don't want that because
    e.g. in the event that we asked the user for a list of strings but the
    user only had one string and just gave us the string itself (e.g. `"abc"`)
    instead of putting it into a list (i.e. `["abc"]`), we don't want to
    iterate over `"abc"` and end up with `["a", "b", "c"]`. Instead, we would
    put `"abc"` into a list on the user's behalf. Note that if `["a", "b", "c"]`
    is really what the user wanted, the user can alternatively specify it
    simply as `list("abc")`."""
    if isinstance(obj, str):
        yield obj
        return
    try:
        yield from iter(obj)
    except TypeError:
        yield obj


def _series_uuid(length: int) -> "pd.Series[str]":
    return pd.Series(map(lambda _: uuid4(), range(length)))


def _raise_if_too_many_dataframes(given: int) -> None:
    limit = len(InferencesRole)
    if not 0 < given <= limit:
        raise ValueError(f"expected between 1 to {limit} dataframes, but {given} were given")


def _coerce_str(obj: Optional[str]) -> str:
    return "" if obj is None else str(obj)


def _coerce_tuple(
    dataframes: Iterable[Union[pd.DataFrame, tuple[Name, pd.DataFrame]]],
) -> Iterator[tuple[Name, pd.DataFrame]]:
    for dataframe in dataframes:
        if isinstance(dataframe, pd.DataFrame):
            yield (_rand_str(), dataframe)
        elif _is_named_df(dataframe):
            yield dataframe
        else:
            raise ValueError(f"unexpected type: {type(dataframe)}")


def _coerce_str_column_names(
    dataframes: Iterable[pd.DataFrame],
) -> Iterator[pd.DataFrame]:
    return (
        df.set_axis(
            df.columns.astype(str),
            axis=1,
        )
        for df in dataframes
    )


_T = TypeVar("_T")


def _add_padding(
    iterable: Iterable[_T],
    padding: Callable[[], _T],
) -> Iterator[_T]:
    return chain(
        (item if item is not None else padding() for item in iterable),
        repeat(padding()),
    )


def _normalize_timestamps(timestamps: "pd.Series[Any]") -> "pd.Series[Any]":
    if timestamps.empty:
        return timestamps
    data_type = timestamps.dtype
    if is_numeric_dtype(data_type) and not is_bool_dtype(data_type):
        return pd.to_datetime(timestamps, unit="s", utc=True)
    if is_datetime64tz_dtype(data_type):
        return timestamps.dt.tz_convert(timezone.utc)
    if is_datetime64_any_dtype(data_type):
        return timestamps.dt.tz_localize(timezone.utc)
    raise ValueError(
        f"When provided, the timestamps must be numeric or datetime, but found {data_type} instead."
    )


_id_pat = re.compile(r"\bid\b", re.IGNORECASE)


def _title_case_no_underscore(name: str) -> str:
    """E.g. `PREDICTION_ID` turns into `Prediction ID`"""
    return _id_pat.sub("ID", name.replace("_", " ").title())


def _jsonify(obj: Any) -> Any:
    if getattr(obj, "__dataclass_fields__", None):
        return {
            attribute.name: _jsonify(value)
            for attribute in fields(obj)
            if attribute.init
            and (value := getattr(obj, attribute.name)) not in ("", None)
            and (not hasattr(value, "__len__") or len(value))
        }
    if isinstance(obj, str):
        return obj
    if isinstance(obj, Iterable):
        return list(map(_jsonify, iter(obj)))
    if isinstance(obj, float) and math.isnan(obj):
        return None
    return obj


def _objectify(json_data: Any) -> Any:
    if isinstance(json_data, str):
        return json_data
    if isinstance(json_data, list):
        return list(map(_objectify, json_data))
    assert isinstance(json_data, dict)
    json_data = {key: _objectify(value) for key, value in json_data.items()}
    # Note that this looks only at the immediate subclasses.
    for cls in CompositeDimensionSpec.__subclasses__():
        try:
            return cls(**json_data)  # type: ignore
        except TypeError:
            pass
    raise ValueError(f"invalid json data: {repr(json_data)}")
    raise ValueError(f"invalid json data: {repr(json_data)}")
