# Suggest running as: WANDB_BASE_URL=http://api.wandb.test python artifact_object_reference_test.py                                                               timothysweeney@Timothys-MacBook-Pro
import shutil

import os
import binascii
import base64
import time
from math import sin, cos, pi
import numpy as np
import sys
from bokeh.plotting import figure

PY3 = sys.version_info.major == 3 and sys.version_info.minor >= 6
if PY3:
    from wandb.sdk.interface import artifacts
else:
    from wandb.sdk_py27.interface import artifacts


WANDB_PROJECT_ENV = os.environ.get("WANDB_PROJECT")
if WANDB_PROJECT_ENV is None:
    WANDB_PROJECT = "test__" + str(round(time.time()) % 1000000)
else:
    WANDB_PROJECT = WANDB_PROJECT_ENV
os.environ["WANDB_PROJECT"] = WANDB_PROJECT

WANDB_SILENT_ENV = os.environ.get("WANDB_SILENT")
if WANDB_SILENT_ENV is None:
    WANDB_SILENT = "true"
else:
    WANDB_SILENT = WANDB_SILENT_ENV
os.environ["WANDB_SILENT"] = WANDB_SILENT

import wandb

columns = ["id", "class_id", "string", "bool", "int", "float", "Image", "Clouds", "HTML", "Video", "Bokeh", "Audio", "np_data"]

def _make_wandb_image(suffix=""):
    class_labels = {1: "tree", 2: "car", 3: "road"}
    test_folder = os.path.dirname(os.path.realpath(__file__))
    im_path = os.path.join(test_folder, "..", "assets", "test{}.png".format(suffix))
    return wandb.Image(
        im_path,
        classes=wandb.Classes([
        {"id": 1, "name": "tree"},
        {"id": 2, "name": "car"},
        {"id": 3, "name": "road"},
    ]),
        boxes={
            "predictions": {
                "box_data": [
                    {
                        "position": {
                            "minX": 0.1,
                            "maxX": 0.2,
                            "minY": 0.3,
                            "maxY": 0.4,
                        },
                        "class_id": 1,
                        "box_caption": "minMax(pixel)",
                        "scores": {"acc": 0.1, "loss": 1.2},
                    },
                    {
                        "position": {
                            "minX": 0.1,
                            "maxX": 0.2,
                            "minY": 0.3,
                            "maxY": 0.4,
                        },
                        "class_id": 2,
                        "box_caption": "minMax(pixel)",
                        "scores": {"acc": 0.1, "loss": 1.2},
                    },
                ],
                "class_labels": class_labels,
            },
            "ground_truth": {
                "box_data": [
                    {
                        "position": {
                            "minX": 0.1,
                            "maxX": 0.2,
                            "minY": 0.3,
                            "maxY": 0.4,
                        },
                        "class_id": 1,
                        "box_caption": "minMax(pixel)",
                        "scores": {"acc": 0.1, "loss": 1.2},
                    },
                    {
                        "position": {
                            "minX": 0.1,
                            "maxX": 0.2,
                            "minY": 0.3,
                            "maxY": 0.4,
                        },
                        "class_id": 2,
                        "box_caption": "minMax(pixel)",
                        "scores": {"acc": 0.1, "loss": 1.2},
                    },
                ],
                "class_labels": class_labels,
            },
        },
        masks={
            "predictions": {
                "mask_data": np.random.randint(0, 4, size=(30, 30)),
                "class_labels": class_labels,
            },
            "ground_truth": {"path": im_path, "class_labels": class_labels},
        },
    )

def _make_point_cloud():
    # Generate a symetric pattern
    POINT_COUNT = 20000

    # Choose a random sample
    theta_chi = pi * np.random.rand(POINT_COUNT, 2)

    def gen_point(theta, chi, i):
        p = sin(theta) * 4.5 * sin(i + 1 / 2 * (i * i + 2)) + \
            cos(chi) * 7 * sin((2 * i - 4) / 2 * (i + 2))

        x = p * sin(chi) * cos(theta)
        y = p * sin(chi) * sin(theta)
        z = p * cos(chi)

        r = sin(theta) * 120 + 120
        g = sin(x) * 120 + 120
        b = cos(y) * 120 + 120

        return [x, y, z, r, g, b]

    def wave_pattern(i):
        return np.array([gen_point(theta, chi, i) for [theta, chi] in theta_chi])

    return wandb.Object3D(wave_pattern(0))

# static assets for comparisons
pc1 = _make_point_cloud()
pc2 = _make_point_cloud()
pc3 = _make_point_cloud()
pc4 = _make_point_cloud()

def _make_bokeh():
    x = [1, 2, 3, 4, 5]
    y = [6, 7, 2, 4, 5]
    p = figure(title="simple line example", x_axis_label='x', y_axis_label='y')
    p.line(x, y, legend_label="Temp.", line_width=2)

    return wandb.data_types.Bokeh(p)

b1 = _make_bokeh()
b2 = _make_bokeh()
b3 = _make_bokeh()
b4 = _make_bokeh()

def _make_html():
    return wandb.Html("<p>Embedded</p><iframe src='https://wandb.ai'></iframe>")

def _make_video():
    return wandb.Video(np.random.randint(0, high=255, size=(4, 1, 10, 10), dtype=np.uint8)) # 1 second video of 10x10 pixels

vid1 = _make_video()
vid2 = _make_video()
vid3 = _make_video()
vid4 = _make_video()

def _make_wandb_audio(frequency, caption):
    SAMPLE_RATE = 44100
    DURATION_SECONDS = 1

    data = np.sin(
        2 * np.pi * np.arange(SAMPLE_RATE * DURATION_SECONDS) * frequency / SAMPLE_RATE
    )
    return wandb.Audio(data, SAMPLE_RATE, caption)

aud1 = _make_wandb_audio(440, "four forty")

aud_ref_https = wandb.Audio(
    "https://wandb-artifacts-refs-public-test.s3-us-west-2.amazonaws.com/StarWars3.wav",
    caption="star wars https"
)
aud_ref_s3 = wandb.Audio(
    "s3://wandb-artifacts-refs-public-test/StarWars3.wav",
    caption="star wars s3"
)
aud_ref_gs = wandb.Audio(
    "gs://wandb-artifact-refs-public-test/StarWars3.wav",
    caption="star wars gs"
)

np_data = np.random.randint(255, size=(4, 16, 16, 3))

def _make_wandb_table():
    classes = wandb.Classes([
        {"id": 1, "name": "tree"},
        {"id": 2, "name": "car"},
        {"id": 3, "name": "road"},
    ])
    table = wandb.Table(
        columns=[c for c in columns[:-1]],
        data=[
            [1, 1, "string1", True, 1, 1.1, _make_wandb_image(), pc1, _make_html(), vid1, b1, aud1],
            [2, 2, "string2", True, 1, 1.2, _make_wandb_image(), pc2, _make_html(), vid2, b2, aud_ref_https],
            [3, 1, "string3", False, -0, -1.3, _make_wandb_image("2"), pc3, _make_html(), vid3, b3, aud_ref_s3],
            [4, 3, "string4", False, -0, -1.4, _make_wandb_image("2"), pc4, _make_html(), vid4, b4, aud_ref_gs],
        ],
    )
    table.cast("class_id", classes.get_type())
    table.add_column(columns[-1], np_data)
    return table

def _make_wandb_joinedtable():
    return wandb.JoinedTable(_make_wandb_table(), _make_wandb_table(), "id")

def _b64_to_hex_id(id_string):
    return binascii.hexlify(base64.standard_b64decode(str(id_string))).decode("utf-8")


# Artifact1.add_reference(artifact_URL) => recursive reference
def test_artifact_add_reference_via_url():
    """ This test creates three artifacts. The middle artifact references the first artifact's file,
    and the last artifact references the middle artifact's reference. The end result of downloading
    the last artifact in a fresh, forth run, should be that all 3 artifacts are downloaded and that
    the file in the last artifact is actually a symlink to the first artifact's file.
    """
    upstream_artifact_name = "upstream_artifact"
    middle_artifact_name = "middle_artifact"
    downstream_artifact_name = "downstream_artifact"

    upstream_local_path = "upstream/local/path/"

    upstream_artifact_path = "upstream/artifact/path/"
    middle_artifact_path = "middle/artifact/path/"
    downstream_artifact_path = "downstream/artifact/path/"

    upstream_local_file_path = upstream_local_path + "file.txt"
    upstream_artifact_file_path = upstream_artifact_path + "file.txt"
    middle_artifact_file_path = middle_artifact_path + "file.txt"
    downstream_artifact_file_path = downstream_artifact_path + "file.txt"

    file_text = "Luke, I am your Father!!!!!"
    # Create a super important file
    if not os.path.exists(upstream_local_path):
        os.makedirs(upstream_local_path)
    with open(upstream_local_file_path, "w") as file:
        file.write(file_text)

    # Create an artifact with such file stored
    with wandb.init() as run:
        artifact = wandb.Artifact(upstream_artifact_name, "database")
        artifact.add_file(upstream_local_file_path, upstream_artifact_file_path)
        run.log_artifact(artifact)

    # Create an middle artifact with such file referenced (notice no need to download)
    with wandb.init() as run:
        artifact = wandb.Artifact(middle_artifact_name, "database")
        upstream_artifact = run.use_artifact(upstream_artifact_name + ":latest")
        artifact.add_reference(
            "wandb-artifact://{}/{}".format(_b64_to_hex_id(upstream_artifact.id),str(upstream_artifact_file_path)),
            middle_artifact_file_path,
        )
        run.log_artifact(artifact)

    # Create a downstream artifact that is referencing the middle's reference
    with wandb.init() as run:
        artifact = wandb.Artifact(downstream_artifact_name, "database")
        middle_artifact = run.use_artifact(middle_artifact_name + ":latest")
        artifact.add_reference(
            "wandb-artifact://{}/{}".format(_b64_to_hex_id(middle_artifact.id),str(middle_artifact_file_path)),
            downstream_artifact_file_path,
        )
        run.log_artifact(artifact)

    # Remove the directories for good measure
    if os.path.isdir("upstream"):
        shutil.rmtree("upstream")
    if os.path.isdir("artifacts"):
        shutil.rmtree("artifacts")

    # Finally, use the artifact (download it) and enforce that the file is where we want it!
    with wandb.init() as run:
        downstream_artifact = run.use_artifact(downstream_artifact_name + ":latest")
        downstream_path = downstream_artifact.download()
        with open(
            os.path.join(downstream_path, downstream_artifact_file_path), "r"
        ) as file:
            assert file.read() == file_text


# # Artifact1.add_reference(artifact2.get_path(file_name))
def test_add_reference_via_artifact_entry():
    """This test is the same as test_artifact_add_reference_via_url, but rather
    than passing the direct URL, we pass an Artifact entry, which will automatically
    resolve to the correct URL
    """
    upstream_artifact_name = "upstream_artifact"
    middle_artifact_name = "middle_artifact"
    downstream_artifact_name = "downstream_artifact"

    upstream_local_path = "upstream/local/path/"

    upstream_artifact_path = "upstream/artifact/path/"
    middle_artifact_path = "middle/artifact/path/"
    downstream_artifact_path = "downstream/artifact/path/"

    upstream_local_file_path = upstream_local_path + "file.txt"
    upstream_artifact_file_path = upstream_artifact_path + "file.txt"
    middle_artifact_file_path = middle_artifact_path + "file.txt"
    downstream_artifact_file_path = downstream_artifact_path + "file.txt"

    file_text = "Luke, I am your Father!!!!!"
    # Create a super important file
    if not os.path.exists(upstream_local_path):
        os.makedirs(upstream_local_path)
    with open(upstream_local_file_path, "w") as file:
        file.write(file_text)

    # Create an artifact with such file stored
    with wandb.init() as run:
        artifact = wandb.Artifact(upstream_artifact_name, "database")
        artifact.add_file(upstream_local_file_path, upstream_artifact_file_path)
        run.log_artifact(artifact)

    # Create an middle artifact with such file referenced (notice no need to download)
    with wandb.init() as run:
        artifact = wandb.Artifact(middle_artifact_name, "database")
        upstream_artifact = run.use_artifact(upstream_artifact_name + ":latest")
        artifact.add_reference(
            upstream_artifact.get_path(upstream_artifact_file_path),
            middle_artifact_file_path,
        )
        run.log_artifact(artifact)

    # Create a downstream artifact that is referencing the middle's reference
    with wandb.init() as run:
        artifact = wandb.Artifact(downstream_artifact_name, "database")
        middle_artifact = run.use_artifact(middle_artifact_name + ":latest")
        artifact.add_reference(
            middle_artifact.get_path(middle_artifact_file_path),
            downstream_artifact_file_path,
        )
        run.log_artifact(artifact)

    # Remove the directories for good measure
    if os.path.isdir("upstream"):
        shutil.rmtree("upstream")
    if os.path.isdir("artifacts"):
        shutil.rmtree("artifacts")

    # Finally, use the artifact (download it) and enforce that the file is where we want it!
    with wandb.init() as run:
        downstream_artifact = run.use_artifact(downstream_artifact_name + ":latest")
        downstream_path = downstream_artifact.download()
        downstream_path = downstream_artifact.download() # should not fail on second download.
        # assert os.path.islink(
        #     os.path.join(downstream_path, downstream_artifact_file_path)
        # )
        with open(
            os.path.join(downstream_path, downstream_artifact_file_path), "r"
        ) as file:
            assert file.read() == file_text

# # Artifact1.get(MEDIA_NAME) => media obj
def test_get_artifact_obj_by_name():
    """Tests tests the ability to instantiate a wandb Media object when passed
    the name of such object. This is the logical inverse of Artifact.add(name).
    TODO: test more robustly for every Media type, nested objects (eg. Table -> Image),
    and references
    """

    with wandb.init() as run:
        artifact = wandb.Artifact("A2", "database")
        image = _make_wandb_image()
        table = _make_wandb_table()
        artifact.add(image, "I1")
        artifact.add(table, "T1")
        run.log_artifact(artifact)

    with wandb.init() as run:
        artifact = run.use_artifact("A2:latest")
        actual_image = artifact.get("I1")
        assert actual_image == image

        actual_table = artifact.get("T1")
        assert actual_table.columns == columns
        assert actual_table.data[0][columns.index("Image")] == image
        assert actual_table.data[1][columns.index("Image")] == _make_wandb_image("2")