"""Shape-first compute, traffic, and activation estimates for batched matmul.

All hardware values and workload dimensions in the executable example are
explicit illustrative inputs. The estimator produces a lower-bound model, not
a latency forecast or benchmark. It deliberately rejects implicit broadcasting
and unspecified traffic ownership.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import date
from decimal import Decimal
import hashlib
import json
import re
from typing import Any


MAX_COUNT = 2**63 - 1
_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/@-]{0,127}$")
_DTYPE_BYTES = {
    "float16": 2,
    "bfloat16": 2,
    "float32": 4,
    "float64": 8,
}
_SAFE_ACCUMULATION_DTYPES = {
    "float16": frozenset(("float16", "float32", "float64")),
    "bfloat16": frozenset(("bfloat16", "float32", "float64")),
    "float32": frozenset(("float32", "float64")),
    "float64": frozenset(("float64",)),
}


def _require_identifier(name: str, value: str) -> None:
    if not isinstance(value, str) or not _IDENTIFIER.fullmatch(value):
        raise ValueError(f"{name} must be a non-empty stable identifier")


def _require_owner(name: str, value: str) -> None:
    _require_identifier(name, value)
    if not value.startswith("team:"):
        raise ValueError(f"{name} must name an accountable team: owner")


def _positive_int(name: str, value: int) -> None:
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive integer")
    if value > MAX_COUNT:
        raise OverflowError(f"{name} exceeds the supported count range")


def _non_negative_int(name: str, value: int) -> None:
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        raise ValueError(f"{name} must be a non-negative integer")
    if value > MAX_COUNT:
        raise OverflowError(f"{name} exceeds the supported count range")


def _checked_product(name: str, factors: tuple[int, ...]) -> int:
    result = 1
    for factor in factors:
        _positive_int(name, factor)
        if result > MAX_COUNT // factor:
            raise OverflowError(f"{name} exceeds the supported count range")
        result *= factor
    return result


def _checked_sum(name: str, values: tuple[int, ...]) -> int:
    result = 0
    for value in values:
        _non_negative_int(name, value)
        if result > MAX_COUNT - value:
            raise OverflowError(f"{name} exceeds the supported count range")
        result += value
    return result


def _canonical(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: _canonical(item) for key, item in value.items()}
    if isinstance(value, tuple):
        return [_canonical(item) for item in value]
    return value


def _content_id(prefix: str, payload: dict[str, Any]) -> str:
    encoded = json.dumps(
        _canonical(payload), sort_keys=True, separators=(",", ":"), allow_nan=False
    ).encode("utf-8")
    return f"{prefix}@sha256:{hashlib.sha256(encoded).hexdigest()}"


@dataclass(frozen=True)
class TensorSpec:
    name: str
    shape: tuple[int, ...]
    dtype: str

    def __post_init__(self) -> None:
        _require_identifier("tensor name", self.name)
        if not isinstance(self.shape, tuple):
            raise TypeError("tensor shape must be an immutable tuple")
        if not 1 <= len(self.shape) <= 8:
            raise ValueError("tensor rank must be between 1 and 8")
        for index, dimension in enumerate(self.shape):
            _positive_int(f"shape[{index}]", dimension)
        if self.dtype not in _DTYPE_BYTES:
            raise ValueError(f"unsupported tensor dtype: {self.dtype}")
        _checked_product("tensor elements", self.shape)

    @property
    def elements(self) -> int:
        return _checked_product("tensor elements", self.shape)

    @property
    def nbytes(self) -> int:
        return _checked_product(
            "tensor bytes", (self.elements, _DTYPE_BYTES[self.dtype])
        )


@dataclass(frozen=True)
class HardwareEnvelope:
    envelope_version: str
    operation_class: str
    operand_dtype: str
    accumulation_dtype: str
    flop_convention: str
    traffic_boundary: str
    peak_flops_per_second: int
    sustained_bytes_per_second: int
    evidence_source_id: str
    measured_at: str
    owner: str

    def __post_init__(self) -> None:
        _require_identifier("envelope_version", self.envelope_version)
        _require_identifier("operation_class", self.operation_class)
        if self.operand_dtype not in _DTYPE_BYTES:
            raise ValueError("unsupported hardware operand_dtype")
        if self.accumulation_dtype not in _DTYPE_BYTES:
            raise ValueError("unsupported hardware accumulation_dtype")
        if (
            self.accumulation_dtype
            not in _SAFE_ACCUMULATION_DTYPES[self.operand_dtype]
        ):
            raise ValueError(
                "hardware accumulation_dtype is incompatible with operand_dtype"
            )
        if self.flop_convention != "fma-is-two-flops":
            raise ValueError("hardware flop_convention must be fma-is-two-flops")
        _require_identifier("hardware traffic_boundary", self.traffic_boundary)
        _positive_int("peak_flops_per_second", self.peak_flops_per_second)
        _positive_int(
            "sustained_bytes_per_second", self.sustained_bytes_per_second
        )
        _require_identifier("evidence_source_id", self.evidence_source_id)
        try:
            parsed_date = date.fromisoformat(self.measured_at)
        except (TypeError, ValueError) as error:
            raise ValueError("measured_at must be an ISO calendar date") from error
        if parsed_date.isoformat() != self.measured_at:
            raise ValueError("measured_at must be an ISO calendar date")
        _require_owner("hardware owner", self.owner)


@dataclass(frozen=True)
class MatmulCostPlan:
    plan_version: str
    operation_id: str
    operation_class: str
    lhs: TensorSpec
    rhs: TensorSpec
    output_name: str
    output_dtype: str
    accumulation_dtype: str
    execution_count: int
    lhs_full_reads: int
    rhs_full_reads: int
    output_full_reads: int
    output_full_writes: int
    execution_liveness: str
    retained_output_copies: int
    additional_retained_activations: tuple[TensorSpec, ...]
    broadcast_policy: str
    flop_convention: str
    traffic_boundary: str
    traffic_assumption_id: str
    shape_owner: str
    kernel_owner: str
    traffic_owner: str
    hardware: HardwareEnvelope

    def __post_init__(self) -> None:
        _require_identifier("plan_version", self.plan_version)
        _require_identifier("operation_id", self.operation_id)
        if self.operation_class != "dense-batched-matmul":
            raise ValueError("operation_class must be dense-batched-matmul")
        if type(self.lhs) is not TensorSpec or type(self.rhs) is not TensorSpec:
            raise TypeError("lhs and rhs must be concrete TensorSpec values")
        if type(self.hardware) is not HardwareEnvelope:
            raise TypeError("hardware must be a concrete HardwareEnvelope")
        _require_identifier("output_name", self.output_name)
        if self.lhs.name == self.rhs.name:
            raise ValueError("lhs and rhs tensor names must be unique")
        if self.output_name in (self.lhs.name, self.rhs.name):
            raise ValueError("output_name must be distinct from input tensor names")
        if len(self.lhs.shape) < 2 or len(self.rhs.shape) < 2:
            raise ValueError("matmul inputs must each have rank at least two")
        if self.lhs.shape[:-2] != self.rhs.shape[:-2]:
            raise ValueError(
                "batch shapes must match exactly; implicit broadcasting is rejected"
            )
        if self.lhs.shape[-1] != self.rhs.shape[-2]:
            raise ValueError("matmul contraction dimensions do not match")
        if self.lhs.dtype != self.rhs.dtype:
            raise ValueError("lhs and rhs dtypes must match")
        if self.output_dtype not in _DTYPE_BYTES:
            raise ValueError("unsupported output_dtype")
        if self.accumulation_dtype not in _DTYPE_BYTES:
            raise ValueError("unsupported accumulation_dtype")
        if self.accumulation_dtype not in _SAFE_ACCUMULATION_DTYPES[self.lhs.dtype]:
            raise ValueError(
                "accumulation_dtype is incompatible with the input dtype"
            )
        _positive_int("execution_count", self.execution_count)
        _positive_int("lhs_full_reads", self.lhs_full_reads)
        _positive_int("rhs_full_reads", self.rhs_full_reads)
        _non_negative_int("output_full_reads", self.output_full_reads)
        _positive_int("output_full_writes", self.output_full_writes)
        _non_negative_int("retained_output_copies", self.retained_output_copies)
        if self.retained_output_copies > self.execution_count:
            raise ValueError("retained_output_copies cannot exceed execution_count")
        if self.execution_liveness == "not-retained":
            if self.retained_output_copies != 0:
                raise ValueError(
                    "not-retained execution_liveness requires zero retained copies"
                )
        elif self.execution_liveness == "serial-reuse":
            if self.retained_output_copies != 1:
                raise ValueError(
                    "serial-reuse execution_liveness requires one retained copy"
                )
        elif self.execution_liveness == "all-executions-live":
            if self.retained_output_copies != self.execution_count:
                raise ValueError(
                    "all-executions-live requires one retained copy per execution"
                )
        else:
            raise ValueError(
                "execution_liveness must be not-retained, serial-reuse, or all-executions-live"
            )
        if not isinstance(self.additional_retained_activations, tuple):
            raise TypeError("additional_retained_activations must be a tuple")
        names = {self.lhs.name, self.rhs.name, self.output_name}
        for activation in self.additional_retained_activations:
            if type(activation) is not TensorSpec:
                raise TypeError(
                    "retained activations must be concrete TensorSpec values"
                )
            if activation.name in names:
                raise ValueError("tensor and activation names must be unique")
            names.add(activation.name)
        if self.broadcast_policy != "exact-batch-no-broadcast":
            raise ValueError("broadcast_policy must be exact-batch-no-broadcast")
        if self.flop_convention != "fma-is-two-flops":
            raise ValueError("flop_convention must be fma-is-two-flops")
        _require_identifier("traffic_boundary", self.traffic_boundary)
        if self.hardware.operation_class != self.operation_class:
            raise ValueError("hardware operation_class does not match the plan")
        if self.hardware.operand_dtype != self.lhs.dtype:
            raise ValueError("hardware operand_dtype does not match the plan")
        if self.hardware.accumulation_dtype != self.accumulation_dtype:
            raise ValueError("hardware accumulation_dtype does not match the plan")
        if self.hardware.flop_convention != self.flop_convention:
            raise ValueError("hardware flop_convention does not match the plan")
        if self.hardware.traffic_boundary != self.traffic_boundary:
            raise ValueError("hardware traffic_boundary does not match the plan")
        _require_identifier("traffic_assumption_id", self.traffic_assumption_id)
        for name in ("shape_owner", "kernel_owner", "traffic_owner"):
            _require_owner(name, getattr(self, name))
        # Validate all derived counts during construction so invalid plans never
        # gain an identity or reach a reporting boundary.
        _derived_counts(self)

    @property
    def plan_id(self) -> str:
        return _content_id(self.plan_version, asdict(self))


@dataclass(frozen=True)
class TensorCostEstimate:
    plan_id: str
    operation_id: str
    output_shape: tuple[int, ...]
    execution_count: int
    execution_liveness: str
    retained_output_copies: int
    flops: int
    traffic_bytes: int
    retained_activation_bytes: int
    arithmetic_intensity_flops_per_byte: Decimal
    compute_lower_bound_seconds: Decimal
    traffic_lower_bound_seconds: Decimal
    lower_bound_seconds: Decimal
    bottleneck: str


def _output_shape(plan: MatmulCostPlan) -> tuple[int, ...]:
    return plan.lhs.shape[:-2] + (plan.lhs.shape[-2], plan.rhs.shape[-1])


def _derived_counts(plan: MatmulCostPlan) -> tuple[int, int, int]:
    batch = _checked_product(
        "batch elements", plan.lhs.shape[:-2] if plan.lhs.shape[:-2] else (1,)
    )
    m, k, n = plan.lhs.shape[-2], plan.lhs.shape[-1], plan.rhs.shape[-1]
    flops = _checked_product(
        "FLOP count", (2, plan.execution_count, batch, m, k, n)
    )
    output_elements = _checked_product("output elements", _output_shape(plan))
    lhs_traffic = _checked_product(
        "lhs traffic bytes",
        (
            plan.lhs.nbytes,
            plan.lhs_full_reads,
            plan.execution_count,
        ),
    )
    rhs_traffic = _checked_product(
        "rhs traffic bytes",
        (
            plan.rhs.nbytes,
            plan.rhs_full_reads,
            plan.execution_count,
        ),
    )
    output_bytes = _checked_product(
        "output bytes", (output_elements, _DTYPE_BYTES[plan.output_dtype])
    )
    output_traffic = _checked_product(
        "output traffic bytes",
        (
            output_bytes,
            plan.output_full_reads + plan.output_full_writes,
            plan.execution_count,
        ),
    )
    traffic = _checked_sum(
        "total traffic bytes", (lhs_traffic, rhs_traffic, output_traffic)
    )

    activation_counts: list[int] = [
        activation.nbytes for activation in plan.additional_retained_activations
    ]
    if plan.retained_output_copies:
        activation_counts.append(
            _checked_product(
                "retained output bytes",
                (output_bytes, plan.retained_output_copies),
            )
        )
    retained = _checked_sum("retained activation bytes", tuple(activation_counts))
    return flops, traffic, retained


def estimate_matmul(plan: MatmulCostPlan) -> TensorCostEstimate:
    if type(plan) is not MatmulCostPlan:
        raise TypeError("plan must be a concrete MatmulCostPlan")
    flops, traffic, retained = _derived_counts(plan)
    if traffic <= 0:
        raise ValueError("traffic model must move at least one byte")
    intensity = Decimal(flops) / Decimal(traffic)
    compute_seconds = Decimal(flops) / Decimal(
        plan.hardware.peak_flops_per_second
    )
    traffic_seconds = Decimal(traffic) / Decimal(
        plan.hardware.sustained_bytes_per_second
    )
    if compute_seconds <= 0 or traffic_seconds <= 0:
        raise ArithmeticError("positive work produced a zero lower bound")
    lower_bound = max(compute_seconds, traffic_seconds)
    if compute_seconds > traffic_seconds:
        bottleneck = "COMPUTE"
    elif traffic_seconds > compute_seconds:
        bottleneck = "TRAFFIC"
    else:
        bottleneck = "BALANCED"
    return TensorCostEstimate(
        plan_id=plan.plan_id,
        operation_id=plan.operation_id,
        output_shape=_output_shape(plan),
        execution_count=plan.execution_count,
        execution_liveness=plan.execution_liveness,
        retained_output_copies=plan.retained_output_copies,
        flops=flops,
        traffic_bytes=traffic,
        retained_activation_bytes=retained,
        arithmetic_intensity_flops_per_byte=intensity,
        compute_lower_bound_seconds=compute_seconds,
        traffic_lower_bound_seconds=traffic_seconds,
        lower_bound_seconds=lower_bound,
        bottleneck=bottleneck,
    )


ILLUSTRATIVE_HARDWARE = HardwareEnvelope(
    envelope_version="illustrative-envelope-v1",
    operation_class="dense-batched-matmul",
    operand_dtype="float16",
    accumulation_dtype="float32",
    flop_convention="fma-is-two-flops",
    traffic_boundary="hbm-to-compute",
    peak_flops_per_second=100_000_000_000_000,
    sustained_bytes_per_second=1_000_000_000_000,
    evidence_source_id="illustrative-capacity-inputs-v1",
    measured_at="2026-08-14",
    owner="team:performance-engineering",
)

ILLUSTRATIVE_PLAN = MatmulCostPlan(
    plan_version="tensor-cost-v1",
    operation_id="illustrative-attention-scores",
    operation_class="dense-batched-matmul",
    lhs=TensorSpec("query", (8, 512, 64), "float16"),
    rhs=TensorSpec("key-transposed", (8, 64, 512), "float16"),
    output_name="attention-scores",
    output_dtype="float32",
    accumulation_dtype="float32",
    execution_count=1,
    lhs_full_reads=1,
    rhs_full_reads=1,
    output_full_reads=0,
    output_full_writes=1,
    execution_liveness="all-executions-live",
    retained_output_copies=1,
    additional_retained_activations=(),
    broadcast_policy="exact-batch-no-broadcast",
    flop_convention="fma-is-two-flops",
    traffic_boundary="hbm-to-compute",
    traffic_assumption_id="illustrative-single-pass-materialized-output-v1",
    shape_owner="team:model-architecture",
    kernel_owner="team:runtime-kernels",
    traffic_owner="team:performance-engineering",
    hardware=ILLUSTRATIVE_HARDWARE,
)


def _shape_text(shape: tuple[int, ...]) -> str:
    return "(" + ",".join(str(dimension) for dimension in shape) + ")"


def format_example() -> str:
    result = estimate_matmul(ILLUSTRATIVE_PLAN)
    lower_bound_us = result.lower_bound_seconds * Decimal(1_000_000)
    return "\n".join(
        [
            "example=illustrative_only",
            f"plan_version={ILLUSTRATIVE_PLAN.plan_version}",
            f"operation={result.operation_id}",
            f"lhs_shape={_shape_text(ILLUSTRATIVE_PLAN.lhs.shape)}",
            f"rhs_shape={_shape_text(ILLUSTRATIVE_PLAN.rhs.shape)}",
            f"output_shape={_shape_text(result.output_shape)}",
            f"execution_count={result.execution_count}",
            f"execution_liveness={result.execution_liveness}",
            f"retained_output_copies={result.retained_output_copies}",
            f"traffic_boundary={ILLUSTRATIVE_PLAN.traffic_boundary}",
            f"flops={result.flops}",
            f"traffic_bytes={result.traffic_bytes}",
            f"retained_activation_bytes={result.retained_activation_bytes}",
            "arithmetic_intensity_flops_per_byte="
            f"{result.arithmetic_intensity_flops_per_byte:.3f}",
            f"lower_bound_us={lower_bound_us:.3f}",
            f"bottleneck={result.bottleneck}",
        ]
    )


if __name__ == "__main__":
    print(format_example())