import argparse
import shutil
from pathlib import Path

import pandas as pd

from src.common import ensure_dirs, load_config, load_json, save_json, seeded_rng
from src.data_generator import aggregate_task_truth, generate_event_log, generate_profiles
from src.experiments import make_figures, retry_experiment, stress_experiment, validation_experiment
from src.parameter_estimation import estimate_parameters


def parse_args():
    parser = argparse.ArgumentParser(description="Agent network traffic data generation and experiments")
    parser.add_argument("--config", default="configs/default.yaml")
    parser.add_argument("--output", default="outputs/latest")
    parser.add_argument("--generator", choices=["local", "llm"], default="local")
    parser.add_argument("--steps", default="generate,estimate,experiment")
    parser.add_argument("--no-fallback", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    root = Path(__file__).resolve().parent
    config_path = (root / args.config).resolve() if not Path(args.config).is_absolute() else Path(args.config)
    output_path = (root / args.output).resolve() if not Path(args.output).is_absolute() else Path(args.output)
    cfg = load_config(config_path)
    dirs = ensure_dirs(output_path)
    shutil.copy2(config_path, output_path / "config_used.yaml")
    steps = {x.strip() for x in args.steps.split(",") if x.strip()}
    rng = seeded_rng(int(cfg["seed"]))

    if "generate" in steps:
        profiles, metadata = generate_profiles(
            args.generator, cfg["generation"]["task_types"], cfg["generation"]["llm"],
            allow_fallback=not args.no_fallback,
        )
        save_json({"metadata": metadata, "profiles": profiles}, dirs["data"] / "behavior_profiles.json")
        train = generate_event_log(profiles, int(cfg["generation"]["train_tasks_per_type"]), rng, cfg["network"])
        test = generate_event_log(profiles, int(cfg["generation"]["test_tasks_per_type"]), rng, cfg["network"])
        train.to_csv(dirs["data"] / "train_events.csv", index=False, encoding="utf-8-sig")
        test.to_csv(dirs["data"] / "test_events.csv", index=False, encoding="utf-8-sig")
        aggregate_task_truth(test).to_csv(dirs["data"] / "task_truth.csv", index=False, encoding="utf-8-sig")
        print("[generate] profiles=%d, train_events=%d, test_events=%d" % (len(profiles), len(train), len(test)))

    if "estimate" in steps:
        train = pd.read_csv(dirs["data"] / "train_events.csv")
        params = estimate_parameters(train)
        save_json(params, dirs["parameters"] / "estimated_parameters.json")
        print("[estimate] task_types=%s" % ",".join(params["task_types"].keys()))

    if "experiment" in steps:
        params = load_json(dirs["parameters"] / "estimated_parameters.json")
        truth = pd.read_csv(dirs["data"] / "task_truth.csv")
        predictions, metrics = validation_experiment(truth, params, cfg, rng)
        stress = stress_experiment(params, cfg, rng)
        retry = retry_experiment(params, cfg, rng)
        predictions.to_csv(dirs["results"] / "validation_predictions.csv", index=False, encoding="utf-8-sig")
        metrics.to_csv(dirs["results"] / "validation_metrics.csv", index=False, encoding="utf-8-sig")
        stress.to_csv(dirs["results"] / "stress_results.csv", index=False, encoding="utf-8-sig")
        retry.to_csv(dirs["results"] / "retry_results.csv", index=False, encoding="utf-8-sig")
        make_figures(predictions, metrics, stress, retry, dirs["figures"])
        summary = {
            "best_model_by_internal_bytes_mape": metrics[metrics["metric"] == "internal_bytes"].sort_values("mape_percent").iloc[0].to_dict(),
            "best_model_by_latency_mape": metrics[metrics["metric"] == "latency_ms"].sort_values("mape_percent").iloc[0].to_dict(),
            "max_stress_arrival_rate": float(stress["arrival_rate"].max()),
            "max_observed_p95_delay_ms": float(stress["p95_delay_ms"].max()),
        }
        save_json(summary, dirs["results"] / "summary.json")
        print("[experiment] results written to %s" % output_path)
        print(metrics[metrics["metric"].isin(["internal_bytes", "latency_ms"])][["model", "metric", "mape_percent"]].to_string(index=False))


if __name__ == "__main__":
    main()

