from __future__ import annotations

import sys
import os
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
WORKSPACE = ROOT.parent
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("XDG_CACHE_HOME", "/tmp")
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(WORKSPACE / "pp-scheduler"))

from pp_scheduler import OpType, UnifiedBigMacVPP  # noqa: E402
from pp_simulator import OpDurationSpec, PipelineSimulator  # noqa: E402


def main() -> None:
    scheduler = UnifiedBigMacVPP(
        pp_size=4,
        vpp_size=2,
        num_microbatches=8,
        ve_forward_limit=3,
    )
    scheduler.generate_schedule()

    simulator = PipelineSimulator.from_scheduler(scheduler)
    result = simulator.simulate(
        {
            OpType.FORWARD: OpDurationSpec(mean=1.0, variance=0),
            OpType.BACKWARD: OpDurationSpec(mean=2.0, variance=0),
            OpType.VEFORWARD: OpDurationSpec(mean=0.8, variance=0.01),
            OpType.VEBACKWARD: OpDurationSpec(mean=0.9, variance=0.01),
            OpType.GENFORWARD: OpDurationSpec(mean=0.6, variance=0.01),
            OpType.GENBACKWARD: OpDurationSpec(mean=0.7, variance=0.01),
        },
        seed=7,
    )
    simulator.validate_result(result)

    print("Simulation summary:")
    for key, value in result.summary.items():
        print(f"  {key}: {value}")

    print("\nFirst 12 simulated ops:")
    for op in result.ops[:12]:
        print(
            f"  rank={op.rank} {op.label:<8} "
            f"start={op.start_time:6.2f} end={op.end_time:6.2f} "
            f"wait={op.wait_time:5.2f} deps={len(op.deps)}"
        )

    print("\nFirst op with dependency reasons:")
    for op in result.ops:
        if op.dep_reasons:
            print(f"  {op.id}")
            for dep_id, reasons in op.dep_reasons.items():
                print(f"    <- {dep_id}: {', '.join(reasons)}")
            break

    trace_path = ROOT / "simulate_demo_trace.json"
    result.to_chrome_trace(trace_path)
    print(f"\nChrome trace written to: {trace_path}")

    perfetto_trace_path = ROOT / "simulate_demo_perfetto_trace.json"
    result.to_chrome_trace(perfetto_trace_path, perfetto_compat=True)
    print(f"Perfetto-compatible trace written to: {perfetto_trace_path}")


if __name__ == "__main__":
    main()