from __future__ import annotations

import shutil
import subprocess
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
DATA = ROOT.parent / 'SimplexTasks-12-data'
RESULTS = ROOT / 'results' / 'tables'
OUTPUTS = ROOT / 'outputs' / 'figures'


def stage_inputs() -> None:
    if not DATA.exists():
        raise SystemExit('SimplexTasks-12-data bundle not found next to the code bundle')
    RESULTS.mkdir(parents=True, exist_ok=True)
    OUTPUTS.mkdir(parents=True, exist_ok=True)
    for subdir in ['figure_inputs', 'table_inputs']:
        src_dir = DATA / 'data' / 'summaries' / subdir
        if not src_dir.exists():
            continue
        for src in src_dir.iterdir():
            if src.is_file():
                shutil.copy2(src, RESULTS / src.name)

    d2_dir = DATA / 'data' / 'synthetic' / 'D2'
    if d2_dir.exists():
        release_dir = ROOT / 'release' / 'simplextasks-12' / 'synthetic' / 'd2_pure_scale'
        release_dir.mkdir(parents=True, exist_ok=True)
        np.savez(
            release_dir / 'task.npz',
            U=np.load(d2_dir / 'U.npy'),
            sigma_true=np.load(d2_dir / 'sigma_true.npy'),
        )


def main() -> None:
    stage_inputs()
    commands = [
        [sys.executable, str(ROOT / 'scripts' / 'make_figures.py'), '--results-dir', str(RESULTS), '--output-dir', str(OUTPUTS)],
        [sys.executable, str(ROOT / 'scripts' / 'plot_pbmc_sensitivity.py'), '--output', str(OUTPUTS / 'fig8_pbmc_sensitivity.pdf')],
    ]
    for cmd in commands:
        subprocess.check_call(cmd, cwd=ROOT)


if __name__ == '__main__':
    main()