import time
from vantage6.client import Client

# Initialise the client
client = Client('http://127.0.0.1', 5000, '/api')
client.authenticate('username', 'password')
client.setup_encryption(None)

# Define algorithm input
input_ = {
    'method': 'master',
    'master': True,
    'kwargs': {
        'org_ids': [2, 3],      # organisations to run kmeans
        'k': 3,                 # number of clusters to compute
        'epsilon': 0.05,        # threshold for convergence criterion
        'max_iter': 300,        # maximum number of iterations to perform
        'columns': [            # columns to be used for clustering
            'sepal_length', 'sepal_width', 'petal_length', 'petal_width'
        ],
        'd_init': 'all',        # data nodes to use for centroids initialisation
        'init_method': 'random',# method for centroids initialisation
        'avg_method': 'k-means' # method used to get global centroids
    }
}

# Send the task to the central server
task = client.task.create(
    collaboration=1,
    organizations=[2, 3],
    name='v6-kmeans-py',
    image='ghcr.io/maastrichtu-cds/v6-kmeans-py:latest',
    description='run kmeans',
    input=input_,
    data_format='json'
)

# Retrieve the results
task_info = client.task.get(task['id'], include_results=True)
while not task_info.get('complete'):
    task_info = client.task.get(task['id'], include_results=True)
    time.sleep(1)
result_info = client.result.list(task=task_info['id'])
results = result_info['data'][0]['result']