import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from .simulator import simulate_queue_stress, simulate_task_batch


def _metrics(actual, predicted):
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)
    return {
        "mae": float(np.mean(np.abs(predicted - actual))),
        "rmse": float(np.sqrt(np.mean((predicted - actual) ** 2))),
        "mape_percent": float(np.mean(np.abs(predicted - actual) / np.maximum(np.abs(actual), 1e-9)) * 100.0),
    }


def validation_experiment(truth, params, cfg, rng):
    scenario_truth = truth.groupby("task_type").agg(
        internal_bytes=("internal_bytes", "mean"),
        latency_ms=("latency_ms", "mean"),
        message_count=("message_count", "mean"),
        tool_calls=("tool_calls", "mean"),
        amplification=("amplification", "mean"),
    ).reset_index()
    global_internal = float(scenario_truth["internal_bytes"].mean())
    global_latency = float(scenario_truth["latency_ms"].mean())
    rows = []
    n = int(cfg["experiments"]["validation_tasks_per_scenario"])
    for _, actual in scenario_truth.iterrows():
        task_type = actual["task_type"]
        for model in ("static_mean", "no_context", "full_model"):
            if model == "static_mean":
                prediction = {
                    "internal_bytes": global_internal,
                    "latency_ms": global_latency,
                    "message_count": float(scenario_truth["message_count"].mean()),
                    "tool_calls": float(scenario_truth["tool_calls"].mean()),
                    "amplification": float(scenario_truth["amplification"].mean()),
                }
            else:
                samples = pd.DataFrame(simulate_task_batch(
                    task_type, n, params, rng, max_retries=cfg["simulation"]["max_retries"],
                    include_context=(model == "full_model"),
                ))
                prediction = {key: float(samples[key].mean()) for key in
                              ("internal_bytes", "latency_ms", "message_count", "tool_calls", "amplification")}
            row = {"task_type": task_type, "model": model}
            for metric, pred in prediction.items():
                row["actual_" + metric] = float(actual[metric])
                row["predicted_" + metric] = pred
            rows.append(row)
    predictions = pd.DataFrame(rows)
    metric_rows = []
    for model, group in predictions.groupby("model"):
        for metric in ("internal_bytes", "latency_ms", "message_count", "tool_calls", "amplification"):
            values = _metrics(group["actual_" + metric], group["predicted_" + metric])
            metric_rows.append({"model": model, "metric": metric, **values})
    return predictions, pd.DataFrame(metric_rows)


def stress_experiment(params, cfg, rng):
    task_type = "collaborative_analysis" if "collaborative_analysis" in params["task_types"] else list(params["task_types"])[-1]
    rows = []
    for rate in cfg["experiments"]["arrival_rates_per_second"]:
        rows.append(simulate_queue_stress(
            float(rate), cfg["experiments"]["stress_duration_seconds"],
            cfg["experiments"]["stress_warmup_seconds"], task_type, params, rng,
            max_concurrency=cfg["network"]["max_concurrency"],
            queue_capacity=cfg["network"]["queue_capacity"],
            max_retries=cfg["simulation"]["max_retries"],
        ))
    return pd.DataFrame(rows)


def retry_experiment(params, cfg, rng):
    task_type = "tool_research" if "tool_research" in params["task_types"] else list(params["task_types"])[0]
    rows = []
    for probability in cfg["experiments"]["timeout_probabilities"]:
        for retry_limit in cfg["experiments"]["retry_limits"]:
            samples = pd.DataFrame(simulate_task_batch(
                task_type, 240, params, rng, timeout_probability=float(probability),
                max_retries=int(retry_limit), include_context=True,
            ))
            rows.append({
                "timeout_probability": probability,
                "max_retries": retry_limit,
                "success_rate": float(samples["success"].mean()),
                "mean_retries": float(samples["retries"].mean()),
                "mean_amplification": float(samples["amplification"].mean()),
                "mean_latency_ms": float(samples["latency_ms"].mean()),
            })
    return pd.DataFrame(rows)


def make_figures(predictions, metrics, stress, retry, figure_dir):
    figure_dir = Path(figure_dir)
    plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "Arial Unicode MS", "DejaVu Sans"]
    plt.rcParams["axes.unicode_minus"] = False

    full = predictions[predictions["model"] == "full_model"]
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
    for ax, metric, label in [(axes[0], "internal_bytes", "内部字节数"), (axes[1], "latency_ms", "任务延迟(ms)")]:
        x, y = full["actual_" + metric], full["predicted_" + metric]
        ax.scatter(x, y, s=65, color="#2E74B5")
        lo, hi = min(x.min(), y.min()), max(x.max(), y.max())
        ax.plot([lo, hi], [lo, hi], "--", color="#777777")
        for _, row in full.iterrows():
            ax.annotate(row["task_type"], (row["actual_" + metric], row["predicted_" + metric]), fontsize=8)
        ax.set_xlabel("模拟实测值"); ax.set_ylabel("模型预测值"); ax.set_title(label)
        ax.grid(alpha=0.25)
    fig.tight_layout(); fig.savefig(figure_dir / "prediction_vs_truth.png", dpi=180); plt.close(fig)

    subset = metrics[metrics["metric"].isin(["internal_bytes", "latency_ms"])]
    pivot = subset.pivot(index="model", columns="metric", values="mape_percent")
    pivot = pivot.rename(
        index={"full_model": "完整模型", "no_context": "无上下文模型", "static_mean": "静态均值模型"},
        columns={"internal_bytes": "内部字节数", "latency_ms": "任务延迟"},
    )
    fig, ax = plt.subplots(figsize=(8, 4.6)); pivot.plot(kind="bar", ax=ax, color=["#2E74B5", "#70AD47"])
    ax.set_ylabel("MAPE (%)"); ax.set_xlabel("模型"); ax.set_title("完整模型与基线对比"); ax.grid(axis="y", alpha=0.25)
    ax.legend(title="指标")
    fig.tight_layout(); fig.savefig(figure_dir / "model_comparison.png", dpi=180); plt.close(fig)

    fig, axes = plt.subplots(2, 2, figsize=(11, 7.5))
    axes[0, 0].plot(stress["arrival_rate"], stress["throughput_per_second"], marker="o")
    axes[0, 0].set_title("到达率—吞吐量")
    axes[0, 1].plot(stress["arrival_rate"], stress["p95_delay_ms"], marker="o", color="#C55A11")
    axes[0, 1].set_title("到达率—P95延迟")
    axes[1, 0].plot(stress["arrival_rate"], stress["queue_peak"], marker="o", color="#A5A5A5")
    axes[1, 0].set_title("到达率—峰值队列")
    axes[1, 1].plot(stress["arrival_rate"], stress["drop_rate"], marker="o", color="#C00000")
    axes[1, 1].set_title("到达率—丢弃率")
    for ax in axes.flat: ax.grid(alpha=0.25); ax.set_xlabel("任务/秒")
    fig.tight_layout(); fig.savefig(figure_dir / "stress_curves.png", dpi=180); plt.close(fig)

    pivot = retry.pivot(index="timeout_probability", columns="max_retries", values="mean_amplification")
    fig, ax = plt.subplots(figsize=(8, 5)); image = ax.imshow(pivot.values, aspect="auto", cmap="YlOrRd")
    ax.set_xticks(range(len(pivot.columns)))
    ax.set_xticklabels([str(x) for x in pivot.columns])
    ax.set_yticks(range(len(pivot.index)))
    ax.set_yticklabels([str(x) for x in pivot.index])
    ax.set_xlabel("最大重试次数"); ax.set_ylabel("超时概率"); ax.set_title("重试导致的流量放大")
    for i in range(pivot.shape[0]):
        for j in range(pivot.shape[1]): ax.text(j, i, "%.2f" % pivot.iloc[i, j], ha="center", va="center", fontsize=8)
    fig.colorbar(image, ax=ax, label="平均流量放大系数")
    fig.tight_layout(); fig.savefig(figure_dir / "retry_heatmap.png", dpi=180); plt.close(fig)
