from collections import defaultdict

import numpy as np


def _stats(series, default=1.0):
    values = np.asarray(series, dtype=float)
    values = values[np.isfinite(values)]
    if len(values) == 0:
        return {"mean": default, "std": 0.0, "cv": 0.1, "p50": default, "p95": default, "count": 0}
    mean = float(np.mean(values))
    std = float(np.std(values))
    return {
        "mean": mean,
        "std": std,
        "cv": max(0.05, std / max(mean, 1e-9)),
        "p50": float(np.quantile(values, 0.50)),
        "p95": float(np.quantile(values, 0.95)),
        "count": int(len(values)),
    }


def estimate_parameters(events):
    parameters = {"global": {}, "task_types": {}}
    transitions = events.groupby(["state_before", "state_after"]).size().to_dict()
    totals = events.groupby("state_before").size().to_dict()
    parameters["global"]["transition_probabilities"] = {
        "%s->%s" % key: float((count + 1.0) / (totals[key[0]] + 8.0))
        for key, count in transitions.items()
    }
    parameters["global"]["duration_ms"] = _stats(events["state_duration_ms"])
    parameters["global"]["message_size_bytes"] = _stats(events["message_size_bytes"])

    for task_type, group in events.groupby("task_type"):
        tasks = group.groupby("task_id")
        task_count = max(1, group["task_id"].nunique())
        initial_tool_requests = group[(group["message_type"] == "tool_request") & (group["retry_count"] == 0)]
        agent_requests = group[group["message_type"] == "agent_request"]
        split_tasks = tasks["state_after"].apply(lambda s: (s == "Split").any())
        retry_tasks = tasks["retry_count"].max()
        tool_requests = group[group["message_type"] == "tool_request"]
        success_tasks = tasks["success"].all()
        external_sizes = group[group["is_external"]]["message_size_bytes"]
        request_sizes = group[group["message_type"].isin(["agent_request", "tool_request"])]["message_size_bytes"]
        response_sizes = group[group["message_type"].isin(["agent_response", "tool_response", "final_response"])]["message_size_bytes"]
        # 5-8ms values represent message-send bookkeeping, not reasoning/service time.
        think_durations = group[
            (group["state_duration_ms"] > 20.0)
            & (group["message_type"].isin(["external_task", "tool_response", "final_response"]))
        ]["state_duration_ms"]
        subtask_counts = tasks["message_type"].apply(lambda s: int((s == "agent_request").sum()))
        parameters["task_types"][task_type] = {
            "phase": str(group["task_phase"].mode().iloc[0]),
            "task_count": task_count,
            "tool_probability": float(len(initial_tool_requests) / max(len(agent_requests), 1)),
            "split_probability": float(split_tasks.mean()),
            "mean_subtasks": float(subtask_counts.mean()),
            "retry_probability": float((tool_requests["retry_count"] > 0).sum() / max(len(tool_requests), 1)),
            "mean_retries_if_any": float(retry_tasks[retry_tasks > 0].mean()) if (retry_tasks > 0).any() else 0.0,
            "failure_probability": float((~success_tasks).mean()),
            "think_time_ms": _stats(think_durations, 500.0),
            "external_size_bytes": _stats(external_sizes, 1500.0),
            "request_size_bytes": _stats(request_sizes, 2000.0),
            "response_size_bytes": _stats(response_sizes, 4000.0),
            "message_count_per_task": _stats(tasks.size(), 3.0),
        }
    return parameters
