import subprocess
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]


alphas = [0.1, 0.3, 0.5, 0.7, 0.9]

alg = "proposed"

comput_dist = "normal"

run_script = str(REPO_ROOT / "scripts" / "run.py")

print("Run the proposed method under different values of alpha...")

for alpha_val in alphas:
    print(f"--- RUN: alg={alg}, comput_dist={comput_dist}, alpha={alpha_val} ---")

    command = [
        "python", run_script,
        "--logdir", str(REPO_ROOT / "results" / "raw"),
        "--modeldir", str(REPO_ROOT / "results" / "checkpoints"),
        "--device", "cuda:0",
        "--alg", alg,
        "--model", "cnn",
        "--dataset", "fmnist",
        "--partition", "noniid-labeldir",
        "--in_dim", "50", # for synthetic data only
        "--n_sample", "5000", # for synthetic data only
        "--n_workers", "10",
        "--comput_dist", comput_dist,
        "--tau_mean", "1",
        "--tau_std", "5",
        "--alpha", str(alpha_val),  # convert floating points to string
        "--lr", "0.01",
        "--tot_step", "5000",
        "--test_freq", "10",
        "--s", "10",
        "--init_seed", "42",
    ]

    try:
        subprocess.run(command, check=True, text=True, capture_output=True)
        print(f"--- SUCCEED: alg={alg}, comput_dist={comput_dist}, alpha={alpha_val} ---")
        # print("Stdout:", result.stdout)
        # print("Stderr:", result.stderr)
    except subprocess.CalledProcessError as e:
        print(f"--- FAIL: alg={alg}, comput_dist={comput_dist}, alpha={alpha_val} ---")
        print("Error:", e)
        print("Stdout:", e.stdout)
        print("Stderr:", e.stderr)
    except FileNotFoundError:
        print(f"Error: cannot find {run_script}. Please check the path.")
        break

    print()

print("DONE")