# platform/deskclaw/experiment/models.py
# ═══════════════════════════════════════════════════════════════
# Controlled Variable Evaluation System — Data Models
#
# AgentBlueprint, VariableMatrix, ExperimentConfig, RunResult,
# and ExperimentResult define the core data structures for
# controlled-variable experiments across multiple agents.
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import itertools
from dataclasses import dataclass, field
from typing import Any

# ── Exceptions ──


class ExperimentError(Exception):
    """Base exception for all experiment module errors."""


class BlueprintValidationError(ExperimentError):
    """AgentBlueprint failed validation."""


class MatrixExpansionError(ExperimentError):
    """VariableMatrix expansion failed."""


# ── AgentBlueprint ──


@dataclass(frozen=True)
class AgentBlueprint:
    """
    Immutable specification for one Agent configuration in an experiment.

    A blueprint captures all controllable parameters for an agent:
    which adapter to use, model selection, sampling temperature,
    system prompt, etc.  The experiment pipeline creates a
    :class:`TemplateAdapter` from each blueprint to enforce
    controlled variables during execution.
    """

    agent_name: str
    adapter_name: str  # registered adapter key (e.g. "self", "nanobot-v1")
    model_id: str = ""
    temperature: float = 0.0
    max_tokens: int = 4096
    system_prompt: str = ""
    capabilities: frozenset[str] = field(default_factory=frozenset)
    metadata: tuple[tuple[str, Any], ...] = ()

    def to_dict(self) -> dict[str, Any]:
        return {
            "agent_name": self.agent_name,
            "adapter_name": self.adapter_name,
            "model_id": self.model_id,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "system_prompt": self.system_prompt,
            "capabilities": sorted(self.capabilities),
            "metadata": dict(self.metadata),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> AgentBlueprint:
        """Construct from a JSON-compatible dict."""
        caps = data.get("capabilities", [])
        meta = data.get("metadata", {})
        return cls(
            agent_name=data["agent_name"],
            adapter_name=data["adapter_name"],
            model_id=data.get("model_id", ""),
            temperature=data.get("temperature", 0.0),
            max_tokens=data.get("max_tokens", 4096),
            system_prompt=data.get("system_prompt", ""),
            capabilities=frozenset(caps) if isinstance(caps, list) else frozenset(caps),
            metadata=tuple(sorted(meta.items()))
            if isinstance(meta, dict)
            else tuple(meta),
        )


# ── VariableMatrix ──


@dataclass(frozen=True)
class VariableMatrix:
    """
    Defines which parameters to vary across experiment runs.

    Each key in *variables* maps to a list of values.  ``expand()``
    produces the full Cartesian product — every combination of variable
    values — as a list of flat dicts.

    Example::

        matrix = VariableMatrix(variables={
            "temperature": [0.0, 0.5, 1.0],
            "max_tokens": [1024, 4096],
        })
        # expand() → 6 combinations
    """

    variables: tuple[tuple[str, tuple[Any, ...]], ...] = ()

    @classmethod
    def from_dict(cls, data: dict[str, list[Any]]) -> VariableMatrix:
        """Build from ``{"var_name": [val1, val2, ...]}``."""
        return cls(variables=tuple((k, tuple(v)) for k, v in sorted(data.items())))

    def expand(self) -> list[dict[str, Any]]:
        """
        Cartesian product of all variable values.

        Returns:
            List of dicts, each mapping variable names to a specific value.
            Empty list if no variables defined.

        Raises:
            MatrixExpansionError: If any variable has an empty value list.
        """
        if not self.variables:
            return [{}]

        keys = [k for k, _ in self.variables]
        value_lists = [v for _, v in self.variables]

        for key, vals in self.variables:
            if not vals:
                raise MatrixExpansionError(
                    f"Variable '{key}' has no values — cannot expand matrix."
                )

        return [dict(zip(keys, combo)) for combo in itertools.product(*value_lists)]

    def to_dict(self) -> dict[str, list[Any]]:
        return {k: list(v) for k, v in self.variables}


# ── ExperimentConfig ──


@dataclass
class ExperimentConfig:
    """
    Full configuration for a controlled-variable experiment.

    Combines multiple :class:`AgentBlueprint` instances with a
    :class:`VariableMatrix` and repetition count.
    """

    name: str
    blueprints: list[AgentBlueprint]
    matrix: VariableMatrix = field(default_factory=lambda: VariableMatrix())
    repetitions: int = 1
    output_dir: str = "./experiment_results"
    timeout_s: int = 600

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "blueprints": [b.to_dict() for b in self.blueprints],
            "matrix": self.matrix.to_dict(),
            "repetitions": self.repetitions,
            "output_dir": self.output_dir,
            "timeout_s": self.timeout_s,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ExperimentConfig:
        blueprints = [AgentBlueprint.from_dict(b) for b in data.get("blueprints", [])]
        matrix = VariableMatrix.from_dict(data.get("matrix", {}))
        return cls(
            name=data["name"],
            blueprints=blueprints,
            matrix=matrix,
            repetitions=data.get("repetitions", 1),
            output_dir=data.get("output_dir", "./experiment_results"),
            timeout_s=data.get("timeout_s", 600),
        )

    def total_runs(self) -> int:
        """Total number of individual runs this experiment will execute."""
        return len(self.blueprints) * len(self.matrix.expand()) * self.repetitions


# ── RunResult ──


@dataclass(frozen=True)
class RunResult:
    """
    Immutable result of a single experiment run.

    Each run corresponds to one (blueprint, variable_combo, repetition) tuple.
    """

    blueprint_name: str
    variables: tuple[tuple[str, Any], ...]  # frozen version of dict
    repetition: int
    status: str  # "passed" | "failed" | "error" | "timeout"
    elapsed_ms: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    api_calls: int = 0
    cost_usd: float = 0.0
    final_answer: str = ""
    error: str | None = None
    trace: tuple[tuple[str, Any], ...] = ()  # frozen version of trace dict

    def to_dict(self) -> dict[str, Any]:
        return {
            "blueprint_name": self.blueprint_name,
            "variables": dict(self.variables),
            "repetition": self.repetition,
            "status": self.status,
            "elapsed_ms": self.elapsed_ms,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "api_calls": self.api_calls,
            "cost_usd": round(self.cost_usd, 6),
            "final_answer": self.final_answer,
            "error": self.error,
            "trace": dict(self.trace),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> RunResult:
        variables = tuple(sorted(data.get("variables", {}).items()))
        trace = tuple(sorted(data.get("trace", {}).items()))
        return cls(
            blueprint_name=data["blueprint_name"],
            variables=variables,
            repetition=data.get("repetition", 0),
            status=data.get("status", "error"),
            elapsed_ms=data.get("elapsed_ms", 0),
            input_tokens=data.get("input_tokens", 0),
            output_tokens=data.get("output_tokens", 0),
            api_calls=data.get("api_calls", 0),
            cost_usd=data.get("cost_usd", 0.0),
            final_answer=data.get("final_answer", ""),
            error=data.get("error"),
            trace=trace,
        )


# ── ExperimentResult ──


@dataclass
class ExperimentResult:
    """
    Aggregate result of a full experiment (all blueprints × matrix × reps).
    """

    config: ExperimentConfig
    runs: list[RunResult] = field(default_factory=list)
    started_at: str = ""
    finished_at: str = ""
    experiment_id: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "experiment_id": self.experiment_id,
            "config": self.config.to_dict(),
            "runs": [r.to_dict() for r in self.runs],
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "total_runs": len(self.runs),
            "passed": sum(1 for r in self.runs if r.status == "passed"),
            "failed": sum(1 for r in self.runs if r.status != "passed"),
        }

    def runs_for_blueprint(self, blueprint_name: str) -> list[RunResult]:
        """Filter runs for a specific blueprint."""
        return [r for r in self.runs if r.blueprint_name == blueprint_name]
