"""检查 Hybrid workload 的资源规划和实际执行配置。

Planner 会在分配大型状态数组前比较候选 method。最终 plan 会给出状态表示、Hilbert
维数、内存估算、worker 数、tolerance 和拒绝原因，结果中另有实际采用的配置。
"""

from __future__ import annotations

import json

from cascaqit import (
    AHSProgram,
    AtomRegister,
    Circuit,
    HybridProgram,
    LocalBackend,
    Waveform,
)
from cascaqit.simulators import SimulationOptions


def main() -> None:
    """运行六 site D-A-D workload，并对照规划值与实际应用值。"""
    sites = 6
    analog = AHSProgram(
        AtomRegister.line(count=sites, spacing=7.0),
        program_id="lesson.hybrid.resource.analog",
    ).drive(
        rabi=Waveform.linear(0.1, 0.6, duration=0.1),
        detuning=Waveform.constant(-0.1, duration=0.1),
        phase=0.0,
    )
    program = (
        HybridProgram("lesson.hybrid.resource")
        .digital("prepare", Circuit(sites).h(0).cx(0, 1))
        .analog("evolve", analog)
        .digital("correct", Circuit(sites).rz(0.2, sites - 1))
        .measure_all()
    )
    result = (
        LocalBackend(analog_time_steps=4)
        .run(
            program,
            shots=16,
            seed=105,
            options=SimulationOptions(
                integrator="fixed_step_krylov",
                dtype="complex64",
                max_steps=4,
                workers=2,
            ),
        )
        .result()
    )
    plan = result.metadata["simulation_plan"]
    execution = result.metadata["simulation_execution_config"]

    payload = {
        "track": "hybrid_researcher",
        "level": "expert",
        "lesson": "resource_planning",
        "facts": {
            "logical_sites": plan["logical_sites"],
            "method_selected": plan["method_selected"],
            "candidate_statuses": {
                item["method"]: item["status"] for item in plan["candidates"]
            },
            "hilbert_dimension": plan["hilbert_dimension"],
            "state_bytes": plan["estimate"]["state_bytes"],
            "estimated_peak_bytes": plan["estimate"]["estimated_peak_bytes"],
            "workers_requested": plan["metadata"]["workers_requested"],
            "workers_applied": execution["workers"],
            "estimated_peak_positive": plan["estimate"]["estimated_peak_bytes"] > 0,
            "execution_dtype": execution["dtype"],
            "tolerance_applied": execution["tolerance_applied"],
            "state_bytes_returned": result.metadata["state_bytes_returned"],
        },
        "boundaries": {
            "hardware_execution": False,
            "cloud_execution": False,
            "network_accessed": False,
            "credentials_loaded": False,
        },
    }
    print(json.dumps(payload, sort_keys=True))


if __name__ == "__main__":
    main()