import sys
import minerva
from   minerva.pier     import Pier

########### PREP ############################

def worker(pier, n):
    mach = pier.machine(ami    = "ami-0a3df8365ed416816",  # dask on ubuntu 22.04 x86
                 instance_type = "t3.medium",
                 username      = "ubuntu",
                 name          = f"dask-worker-{n}",
                 variables     = {"type":   "worker",
                                  "number": n},
                 disk_size     = 512)
    return mach

def scheduler(pier):
    mach = pier.machine(ami    = "ami-0a3df8365ed416816",  # dask on ubuntu 22.04 x86
                 instance_type = "t3.medium",
                 username      = "ubuntu",
                 name          = f"dask-scheduler",
                 variables     = {"type": "scheduler"},
                 disk_size     = 32)
    return mach

########## CLUSTER ##########################

m       = minerva.Minerva("hay")
pier    = m.pier(subnet_id = "subnet-05eb26d8649a093e1", # project-subnet-public1-us-east-1a
                 sg_groups = ["sg-0f9e555954e863954",    # ssh
                              "sg-0b34a3f7398076545",    # default
                              "sg-04cd2626d91ac093c"],   # dask (8786, 8787)
                 key_pair  = ("Ari-Brown-HAY", "~/.ssh/Ari-Brown-HAY.pem"),
                 iam       = "Minerva")

cluster = pier.cluster(scheduler, worker, num_workers=int(sys.argv[1]))
cluster.start()
print(cluster.public_location)

print("press enter to terminate cluster")
input()
cluster.terminate()