import numpy as np
from astropy.table import Table
from csst_dag.dag import Dispatcher

# # 1221 plan recs, 36630 data recs
plan_basis, data_basis = Dispatcher.load_test_data()

# # 666 task/s
task_list_via_file = Dispatcher.dispatch_file(plan_basis, data_basis)
t = Table(task_list_via_file)
np.unique(t["success"])
np.unique(t["n_relevant_plan"])
np.unique(t["n_relevant_data"])
print(t["task"])
print(sum(t["success"]))

# # 13 task/s @n_jobs=1, 100*10 task/s @n_jobs=10 (max)
task_list_via_detector = Dispatcher.dispatch_detector(
    plan_basis, data_basis[::10], n_jobs=1
)
t = Table(task_list_via_detector)
np.unique(t["n_relevant_plan"], return_counts=True)
np.unique(t["success"], return_counts=True)

# 16 task/s @n_jobs=1, 130*10 tasks/s @n_jobs=10 (max) 🔼
task_list_via_obsid = Dispatcher.dispatch_obsid(plan_basis, data_basis, n_jobs=10)
t = Table(task_list_via_obsid)
np.unique(t["n_relevant_plan"], return_counts=True)
np.unique(t["success"], return_counts=True)

# 13s/task
task_list_via_obsgroup = Dispatcher.dispatch_obsgroup(plan_basis, data_basis)
t = Table(task_list_via_obsgroup)
print(
    sum(_["success"] for _ in task_list_via_obsgroup),
    "/",
    len(task_list_via_obsgroup),
)