"""
The following is a simple example algorithm.

It is meant to run within a container.

To run the container locally, you can call the following bash script:

  ./do_test_run.sh

This will start the inference and reads from ./test/input and writes to ./test/output

To save the container and prep it for upload to Grand-Challenge.org you can call:

  ./do_save.sh

Any container that shows the same behaviour will do, this is purely an example of how one COULD do it.

Reference the documentation to get details on the runtime environment on the platform:
https://grand-challenge.org/documentation/runtime-environment/

Happy programming!
"""
from pathlib import Path
import json
{% if algorithm_input_sockets | has_image or algorithm_output_sockets | has_image -%}
from glob import glob
import SimpleITK
import numpy
{%- endif %}

INPUT_PATH = Path("/input")
OUTPUT_PATH = Path("/output")
RESOURCE_PATH = Path("resources")

def run():
    # The key is a tuple of the slugs of the input sockets
    interface_key = get_interface_key()

    # Lookup the handler for this particular set of sockets (i.e. the interface)
    handler = {
    {% for interface_name, interface_key in algorithm_interface_names|zip(algorithm_interface_keys) -%}
    {{ interface_key }} : {{ interface_name }}_handler,
    {% endfor %}
    }[interface_key]

    # Call the handler
    return handler()

{% for interface_name, interface in algorithm_interface_names|zip(phase.algorithm_interfaces) -%}
def {{ interface_name }}_handler():
    # Read the input
    {% for socket in interface["inputs"] -%}
    {% set py_slug = socket.slug.replace("-", "_") -%}
    {% if socket | is_image -%}
    input_{{ py_slug }} = load_image_file_as_array(
        location=INPUT_PATH / "{{ socket.relative_path }}",
    )
    {% endif -%}
    {% if socket | is_json -%}
    input_{{ py_slug }} = load_json_file(
         location=INPUT_PATH / "{{ socket.relative_path }}",
    )
    {% endif -%}
    {% if socket | is_file -%}
    input_{{ py_slug }} = load_file(
         location=INPUT_PATH / "{{ socket.relative_path }}",
    )
    {% endif -%}
    {% endfor %}
    # Process the inputs: any way you'd like
    _show_torch_cuda_info()

    # Some additional resources might be required, include these in one of two ways.

    # Option 1: part of the Docker-container image: resources/
    resource_dir = Path("/opt/app/resources")
    with open(resource_dir / "some_resource.txt", "r") as f:
        print(f.read())

    # Option 2: upload them as a separate tarball to Grand Challenge (go to your Algorithm > Models). The resources in the tarball will be extracted to `model_dir` at runtime.
    model_dir = Path("/opt/ml/model")
    with open(model_dir / "a_tarball_subdirectory" / "some_tarball_resource.txt", "r") as f:
        print(f.read())

    # For now, let us make bogus predictions
    {%- for socket in interface["outputs"] %}
    output_{{ socket.slug.replace("-", "_")}} =
        {%- if socket | has_example_value %}
            {%- if socket.example_value is string -%}
                "{{ socket.example_value }}"
            {%- else -%}
                {{ socket.example_value }}
            {%- endif %}
        {%- elif socket | is_image %} numpy.eye(4, 2)
        {%- elif socket | is_json %} {"content": "should match the required format"}
        {%- elif socket | is_file %} "content: should match the required format"
        {% endif %}
    {%- endfor %}

    # Save your output
    {% for socket in interface["outputs"] -%}
    {% set py_slug = socket.slug.replace("-", "_") -%}
    {% if socket | is_image -%}
    write_array_as_image_file(
        location=OUTPUT_PATH / "{{ socket.relative_path }}",
        array=output_{{ py_slug }},
    )
    {% endif -%}
    {% if socket | is_json -%}
    write_json_file(
        location=OUTPUT_PATH / "{{ socket.relative_path }}",
        content=output_{{ py_slug }}
    )
    {% endif -%}
    {% if socket | is_file -%}
    write_file(
        location=OUTPUT_PATH / "{{ socket.relative_path }}",
        content=output_{{ py_slug }}
    )
    {% endif -%}
    {% endfor %}
    return 0
{% endfor %}

def get_interface_key():
    # The inputs.json is a system generated file that contains information about
    # the inputs that interface with the algorithm
    inputs = load_json_file(
        location=INPUT_PATH / "inputs.json",
    )
    socket_slugs = [sv["interface"]["slug"] for sv in inputs]
    return tuple(sorted(socket_slugs))

def load_json_file(*, location):
    # Reads a json file
    with open(location, 'r') as f:
        return json.loads(f.read())

{%- if algorithm_output_sockets| has_json %}
def write_json_file(*, location, content):
    # Writes a json file
    with open(location, 'w') as f:
        f.write(json.dumps(content, indent=4))
{%- endif %}


{%- if algorithm_input_sockets | has_image %}
def load_image_file_as_array(*, location):
    # Use SimpleITK to read a file
    input_files = glob(str(location / "*.tif")) + glob(str(location / "*.tiff")) + glob(str(location / "*.mha"))
    result = SimpleITK.ReadImage(input_files[0])

    # Convert it to a Numpy array
    return SimpleITK.GetArrayFromImage(result)
{%- endif %}


{%- if algorithm_output_sockets | has_image %}
def write_array_as_image_file(*, location, array):
    location.mkdir(parents=True, exist_ok=True)

    # You may need to change the suffix to .tif to match the expected output
    suffix = ".mha"

    image = SimpleITK.GetImageFromArray(array)
    SimpleITK.WriteImage(
        image,
        location / f"output{suffix}",
        useCompression=True,
    )
{%- endif %}


{%- if algorithm_input_sockets | has_file %}
# Note to the developer:
#   the following function is very generic and should likely
#   be adopted to something more specific for your algorithm/challenge
def load_file(*, location):
    # Reads the content of a file
    with open(location) as f:
        return f.read()
{%- endif %}


{%- if algorithm_output_sockets | has_file %}
# Note to the developer:
#   the following function is very generic and should likely
#   be adopted to something more specific for your algorithm/challenge
def write_file(*, location, content):
    # Write the content to a file
    with open(location, 'w') as f:
        return f.write(content)
{%- endif %}


def _show_torch_cuda_info():
    import torch

    print("=+=" * 10)
    print("Collecting Torch CUDA information")
    print(f"Torch CUDA is available: {(available := torch.cuda.is_available())}")
    if available:
        print(f"\tnumber of devices: {torch.cuda.device_count()}")
        print(f"\tcurrent device: { (current_device := torch.cuda.current_device())}")
        print(f"\tproperties: {torch.cuda.get_device_properties(current_device)}")
    print("=+=" * 10)


if __name__ == "__main__":
    raise SystemExit(run())
