"""
The following is a simple example algorithm.

It is meant to run within a container.

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

  ./test_run.sh

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

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

  docker save example-algorithm-{{ cookiecutter.phase.slug }} | gzip -c > example-algorithm-{{ cookiecutter.phase.slug }}.tar.gz

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

Happy programming!
"""
from pathlib import Path
{% if cookiecutter.phase.inputs | has_json or cookiecutter.phase.outputs | has_json -%}
import json
{%- endif %}
{% if cookiecutter.phase.inputs | has_image or cookiecutter.phase.outputs | has_image -%}
from glob import glob
import SimpleITK
import numpy
{%- endif %}

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

{% if "__should_fail" in cookiecutter.phase -%} 1/0 {%- endif %}

def run():
    # Read the input
    {% for ci in cookiecutter.phase.inputs -%}
    {% set py_slug = ci.slug | replace("-", "_") -%}
    {% if ci | is_image -%}
    {{ py_slug }} = load_image_file_as_array(
        location=INPUT_PATH / "{{ ci.relative_path }}",
    )
    {% endif -%}
    {% if ci | is_json -%}
    {{ py_slug }} = load_json_file(
         location=INPUT_PATH / "{{ ci.relative_path }}",
    )
    {% endif -%}
    {% endfor %}
    # Process the inputs: any way you'd like
    _show_torch_cuda_info()

    with open(RESOURCE_PATH / "some_resource.txt", "r") as f:
        print(f.read())

    # For now, let us set make bogus predictions
    {%- for ci in cookiecutter.phase.outputs %}
    {{ ci.slug | replace("-", "_")}} =
        {%- if ci | is_image %} numpy.eye(4, 2)
        {%- elif ci | is_json %} {"content": "should match the required format"}
        {%- elif ci | is_file %} "content: should match the required format"
        {% endif %}
    {%- endfor %}

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


def load_json_file(*, location):
    # Reads a json file
    with open(location, 'r') as f:
        return json.loads(f.read())
{%- endif %}
{%- if cookiecutter.phase.outputs | 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 cookiecutter.phase.inputs | has_image %}


def load_image_file_as_array(*, location):
    # Use SimpleITK to read a file
    input_files = 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 cookiecutter.phase.outputs | 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 .tiff to match the expected output
    suffix = ".mha"

    image = SimpleITK.GetImageFromArray(array)
    SimpleITK.WriteImage(
        image,
        location / f"output{suffix}",
        useCompression=True,
    )
{%- endif %}
{%- if cookiecutter.phase.inputs | has_file %}


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


# Note to the challenge hosts:
#   the following function is very generic and should likely
#   be adopted to something more specific for your 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())
