#!/usr/bin/env python

__doc__ = """

This example shows how to use a client to access a 

remote Sybil server (running Ark) to predict risk scores for a set of DICOM files.

The server must be started separately.

https://github.com/reginabarzilaygroup/Sybil/wiki

https://github.com/reginabarzilaygroup/ark/wiki

"""

import json

import os

import numpy as np

import requests

import sybil.utils.visualization

from utils import get_demo_data

if __name__ == "__main__":

    dicom_files = get_demo_data()

    serie = sybil.Serie(dicom_files)

    # Set the URL of the remote Sybil server

    ark_hostname = "localhost"

    ark_port = 5000

    # Set the URL of the remote Sybil server

    ark_host = f"http://{ark_hostname}:{ark_port}"

    data_dict = {"return_attentions": True}

    payload = {"data": json.dumps(data_dict)}

    # Check if the server is running and reachable

    resp = requests.get(f"{ark_host}/info")

    if resp.status_code != 200:

        raise ValueError(f"Failed to connect to ARK server. Status code: {resp.status_code}")

    info_data = resp.json()["data"]

    assert info_data["modelName"].lower() == "sybil", "The ARK server is not running Sybil"

    print(f"ARK server info: {info_data}")

    # Submit prediction to ARK server.

    files = [('dicom', open(file_path, 'rb')) for file_path in dicom_files]

    r = requests.post(f"{ark_host}/dicom/files", files=files, data=payload)

    _ = [f[1].close() for f in files]

    if r.status_code != 200:

        raise ValueError(f"Error occurred processing DICOM files. Status code: {r.status_code}.\n{r.text}")

    r_json = r.json()

    predictions = r_json["data"]["predictions"]

    scores = predictions[0]

    print(f"Risk scores: {scores}")

    attentions = predictions[1]

    attentions = np.array(attentions)

    print(f"Ark received attention shape: {attentions.shape}")

    # Visualize attention maps

    save_directory = "remote_ark_sybil_attention_output"

    print(f"Writing attention images to {save_directory}")

    images = serie.get_raw_images()

    overlayed_images = sybil.utils.visualization.build_overlayed_images(images, attentions, gain=3)

    if save_directory is not None:

        serie_idx = 0

        save_path = os.path.join(save_directory, f"serie_{serie_idx}")

        sybil.utils.visualization.save_images(overlayed_images, save_path, f"serie_{serie_idx}")

    print(f"Finished writing attention images to {save_directory}")