"""Variable types for rTisanePy.

Ports Unit, Participant, Time, Measure types and their subtypes
from rTisane's S4 class hierarchy (aaa-classes.R).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Union


# ---------------------------------------------------------------------------
# Number-value helpers (for numberOfInstances)
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class Exactly:
    """Exact number of instances."""
    value: int = 1


@dataclass(frozen=True)
class AtMost:
    """Upper-bound number of instances."""
    value: int = 1


@dataclass(frozen=True)
class Per:
    """Ratio: *number* instances per *variable*."""
    number: Union[int, Exactly, AtMost]
    variable: "AbstractVariable"


NumberOfInstances = Union[int, Exactly, AtMost, Per]


# ---------------------------------------------------------------------------
# Variables
# ---------------------------------------------------------------------------

@dataclass
class AbstractVariable:
    """Base for all variables."""
    name: str


@dataclass
class Unit(AbstractVariable):
    """Entity from which data is collected (e.g. participant, subject)."""
    cardinality: int = 0
    nests_within: Union["Unit", list["Unit"], None] = None


@dataclass
class Participant(Unit):
    """Alias for Unit."""
    pass


@dataclass
class Time(AbstractVariable):
    """Time variable (ordinal or nominal)."""
    order: list | None = None
    cardinality: int = 0

    def __post_init__(self):
        if self.cardinality == 0 and not self.order:
            raise ValueError("Provide at least cardinality or order for Time.")
        if self.order and self.cardinality == 0:
            self.cardinality = len(self.order)
        if self.order and self.cardinality != 0 and len(self.order) != self.cardinality:
            raise ValueError("order length must match cardinality.")


# ---------------------------------------------------------------------------
# Measures
# ---------------------------------------------------------------------------

@dataclass
class Measure(AbstractVariable):
    """Base measure class — not instantiated directly."""
    unit: Unit = field(default_factory=lambda: Unit(""))
    number_of_instances: NumberOfInstances = 1


@dataclass
class Continuous(Measure):
    """Continuous measure (scores, temperature, etc.)."""
    pass


@dataclass
class Counts(Measure):
    """Count measure."""
    pass


@dataclass
class Categories(Measure):
    """Base for categorical measures."""
    pass


@dataclass
class UnorderedCategories(Categories):
    """Unordered (nominal) categories."""
    cardinality: int = 2


@dataclass
class OrderedCategories(Categories):
    """Ordered (ordinal) categories."""
    order: list = field(default_factory=list)
    cardinality: int = 0

    def __post_init__(self):
        if self.order and self.cardinality == 0:
            self.cardinality = len(self.order)


@dataclass
class UnobservedVariable(AbstractVariable):
    """Unobserved / latent variable (internal use)."""
    name: str = "Unobserved"


# ---------------------------------------------------------------------------
# Constructor functions (match rTisane's standalone API)
# ---------------------------------------------------------------------------

def continuous(unit: Unit, name: str, number_of_instances: NumberOfInstances = 1) -> Continuous:
    """Create a Continuous measure."""
    return Continuous(name=name, unit=unit, number_of_instances=number_of_instances)


def counts(unit: Unit, name: str, number_of_instances: NumberOfInstances = 1) -> Counts:
    """Create a Counts measure."""
    return Counts(name=name, unit=unit, number_of_instances=number_of_instances)


def categories(
    unit: Unit,
    name: str,
    *,
    cardinality: int | None = None,
    order: list | None = None,
    number_of_instances: NumberOfInstances = 1,
) -> UnorderedCategories | OrderedCategories:
    """Create a categorical measure (ordered if *order* is given, else unordered)."""
    if order is not None:
        return OrderedCategories(
            name=name, unit=unit, number_of_instances=number_of_instances,
            order=order, cardinality=len(order),
        )
    if cardinality is None:
        raise ValueError("Provide either cardinality or order.")
    return UnorderedCategories(
        name=name, unit=unit, number_of_instances=number_of_instances,
        cardinality=cardinality,
    )
