import sys
import minerva
from   minerva.pier     import Pier

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

DASK_BASE = "ami-0399a4f70ca684620" # dask on ubuntu 22.04 x86

def worker(pier, n):
    mach = pier.machine(ami    = DASK_BASE,
                 instance_type = "m5.large",
                 username      = "ubuntu",
                 name          = f"dask-worker-{n}",
                 variables     = {"type":   "worker",
                                  "number": n},
                 disk_size     = 512)
    return mach

def scheduler(pier):
    mach = pier.machine(ami    = DASK_BASE,
                 instance_type = "m5.large", # "r5.xlarge",
                 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()
print(f"dashboard: http://{cluster.scheduler.public_ip}:8787/")
print(f"cluster:   {cluster.public_location}")
print()

print("type `exit()` to terminate the cluster")
print()

import IPython
IPython.embed()

cluster.terminate()