#!/usr/bin/env python3
"""Unit tests for analyze_data.py.

If this test fails because you, for example, updated the output of
'extract_data.{cpp,py}', simply replace the 'test_exp' tar file with a new
extraction. Only include a single row of data. The goal of this test is just to
ensure we can read from extracted files and that we fail properly when data is
not available.

The current 'test_exp.tar.gz' directory structure is like this:

    test_exp/
    |-> shaA/
        |-> <chum URI>.mp4.loc
        |-> <chum URI>.pcp.csv
        |-> <chum URI>.plan.csv
        |-> <chum URI>.pred.csv
    |-> shaB/
        |-> <chum URI>.mp4.loc
        |-> <chum URI>.pcp.csv
        |-> <chum URI>.plan.csv
        |-> <chum URI>.pred.csv
    test_steering_only_exp/
    |-> shaA/
        |-> <steering scenario hash>.steering.json
    |-> shaB/
        |-> <steering scenario hash>.steering.json

To update 'test_exp.tar.gz', find a recent PISCES run:

    https://piper.zooxlabs.com/?name=pisces.

Chum data is cleaned up regularly, so unless the run is recent, you won't find
any data.

Locate the extract stage, e.g.

    https://piper.zooxlabs.com/pipeline/e2fbf9f0-51db-497d-8ec4-07ec57f2f082/stage/extract.extract

If you updated the extract stage, copy-paste the extract command to the terminal, e.g.

    bazel run //mined_metric/builder/metrics_impl/pisces/utils:extract_data -- \
        --id <PISCES ID> \
        -o /mnt/nautilus_rw/pisces/data

If you didn't update the extract stage, there's no need to run the extract stage
again. You can use the existing extraction results by concatenating the id with
the output path, e.g.

    ls /mnt/nautilus_rw/pisces/data/<PISCES ID>

There should be two subdirectories corresponding to the two "releases" that were
compared in the experiment. In the 'test_exp.tar.gz' directory structure, these
subdirectories are renamed to "shaA" and "shaB". Find the same "scenario" (i.e.
chum URI) in both subdirectories and copy them to your local file system.

    cp /mnt/nautilus_rw/pisces/data/<PISCES ID>/<shaA>/<chum URI>* ~/test_exp/shaA/
    cp /mnt/nautilus_rw/pisces/data/<PISCES ID>/<shaB>/<chum URI>* ~/test_exp/shaB/

Also copy a single steering scenario from each "release":

    cp /mnt/nautilus_rw/pisces/data/<PISCES ID>/<shaA>/<steering scenario hash>.steering.json ~/test_steering_only_exp/shaA/
    cp /mnt/nautilus_rw/pisces/data/<PISCES ID>/<shaB>/<steering scenario hash>.steering.json ~/test_steering_only_exp/shaB/

To make this test complete faster in CI, modify all 14 CSV files so they contain
just two rows each: a header row and a single data row. **NOTE** - the extracted
planner data (*.plan.csv) needs to contain the *last* row, not the first,
because planner data is extracted in reverse order.

Then tar up the two directories:

    tar -zcvf test_exp.tar.gz test_exp/ test_steering_only_exp/

You can list the contents of the resulting tar file (to confirm the expected
directory structure, for example) using:

    tar -ztvf test_exp.tar.gz

When ready, upload to S3/Z3 with:

    ./scripts/build/upload-test-data.sh ~/test_exp.tar.gz

Finally, update the :test_exp.tar.gz z3_test_data BUILD target using the sha256
printed to the terminal. Make sure this test passes locally before pushing to
GitHub.
"""
import os
import shutil
import sys
import tarfile
import tempfile

from bokeh.models import Tabs
import pandas as pd
import pytest

import mined_metric.builder.metrics_impl.pisces.errors as errs
from mined_metric.builder.metrics_impl.pisces.utils.analyze_data import (
    build_report_tabs,
    load_experiment,
)


@pytest.fixture
def mock_pisces_results(scope="module"):
    """Mock PISCES results.

    Extracts the test data to a temporary directory. Adds an "experiment"
    (directory) with faulty data to the two existing "experiments" extracted
    from 'test_exp.tar.gz':

    test_exp/
        Contains CSVs for Planner, Prediction, and Perception extracted from a
        "nominal" (non-steering) PISCES test.
    test_steering_only_exp/
        Contains CSVs for Planner, Prediction, and Perception extracted from a
        steering comparison PISCES test.
    test_exp_missing_cols/
        A copy of 'test_exp/' but modified to have missing / faulty data.
    test_full_exp/
        The combined contents of 'test_exp/' and 'test_steering_only_exp/' to
        have a "full" experiment with at least one "nominal" scenario and at
        least one "steering-only" scenario.
    """
    test_data = "mined_metric/builder/metrics_impl/pisces/utils/test_exp.tar.gz"
    tmpdir = tempfile.mkdtemp(dir=os.environ.get("TEST_TMPDIR", None))
    try:
        tar = tarfile.open(test_data, "r:gz")
        tar.extractall(path=tmpdir)
        tar.close()

        # Combine the contents of 'test_exp' and 'test_steering_only_exp' into a
        # new "full" experiment directory.
        shutil.copytree(
            os.path.join(tmpdir, "test_exp"),
            os.path.join(tmpdir, "test_full_exp"),
        )
        for src_dir, _, files in os.walk(
            os.path.join(tmpdir, "test_steering_only_exp")
        ):
            dst_dir = src_dir.replace("test_steering_only_exp", "test_full_exp")
            for file_ in files:
                shutil.copy(os.path.join(src_dir, file_), dst_dir)

        # Make a copy of 'test_full_exp' but modify to have missing / faulty data.
        shutil.copytree(
            os.path.join(tmpdir, "test_full_exp"),
            os.path.join(tmpdir, "test_exp_missing_cols"),
        )

        file_to_change = os.path.join(
            tmpdir,
            "test_exp_missing_cols",
            "shaA",
            "20200115T004406-kitt_22@1579049341480000000-1579049345980000000.plan.csv",
        )

        df_to_modify = pd.read_csv(file_to_change)
        del df_to_modify["decision_vel"]
        df_to_modify.to_csv(file_to_change, index=False)

        yield dict(
            directory=tmpdir,
            experiment_id="test_exp",
            missing_col_experiment_id="test_exp_missing_cols",
            steering_experiment_id="test_steering_only_exp",
            full_experiment_id="test_full_exp",
        )

    finally:
        shutil.rmtree(tmpdir)


def expect_load_ok(directory, experiment_id):
    exp = load_experiment(directory, experiment_id)

    stats_a = exp.describe("shaA")["Planner"]
    stats_b = exp.describe("shaB")["Planner"]

    assert len(stats_a) == len(
        stats_b
    ), "Expected experiments to have same number of scenarios"

    for i in range(len(stats_a)):
        assert (
            stats_a.ticks.iloc[i] == stats_b.ticks.iloc[i]
        ), "Expected tick counts to be equal"

    return exp


@pytest.mark.parametrize(
    "experiment_id",
    ["full_experiment_id", "steering_experiment_id", "experiment_id"],
)
def test_load_ok(mock_pisces_results, experiment_id):
    expect_load_ok(
        mock_pisces_results["directory"], mock_pisces_results[experiment_id]
    )


@pytest.mark.parametrize(
    "experiment_id",
    ["full_experiment_id", "steering_experiment_id", "experiment_id"],
)
def test_compare_ok(mock_pisces_results, experiment_id):
    exp = expect_load_ok(
        mock_pisces_results["directory"], mock_pisces_results[experiment_id]
    )

    comparators = exp.compare("shaA", "shaB")

    # congrats! you successfully compared
    assert comparators is not None, "Expected nontrivial dataframe"


def test_load_fail_no_data(mock_pisces_results):
    with pytest.raises(errs.PiscesNoDataFound):
        load_experiment(
            mock_pisces_results["directory"], "nonexistent_experiment"
        )


def test_compare_fail_bad_format_data(mock_pisces_results):
    exp = expect_load_ok(
        mock_pisces_results["directory"],
        mock_pisces_results["missing_col_experiment_id"],
    )

    with pytest.raises(KeyError):
        comparators = exp.compare("shaA", "shaB")
        for _, comparator in comparators.items():
            [comp.diff() for comp in comparator if comp is not None]


@pytest.mark.parametrize(
    "experiment_id",
    ["full_experiment_id", "steering_experiment_id", "experiment_id"],
)
def test_describe_ok(mock_pisces_results, experiment_id):
    exp = expect_load_ok(
        mock_pisces_results["directory"], mock_pisces_results[experiment_id]
    )
    comparators = exp.compare("shaA", "shaB")

    stats_a = exp.describe("shaA")
    stats_b = exp.describe("shaB")

    assert all(
        [
            key in stats
            for key in ["Planner", "Prediction"]
            for stats in [stats_a, stats_b]
        ]
    )


@pytest.mark.parametrize(
    "experiment_id",
    ["full_experiment_id", "steering_experiment_id", "experiment_id"],
)
def test_build_report_tabs(mock_pisces_results, experiment_id):
    """Tests building all and steering-only report tabs."""
    candidate = "shaA"
    control = "shaB"

    exp = load_experiment(
        mock_pisces_results["directory"], mock_pisces_results[experiment_id]
    )
    comparators = exp.compare(candidate, control)
    candidate_stats = exp.describe(candidate)
    control_stats = exp.describe(control)

    tabs, changed_metrics = build_report_tabs(
        comparators,
        exp,
        control,
        "control_sha",
        control_stats,
        candidate,
        "candidate_sha",
        candidate_stats,
        "argus_host",
        None,
        "",
    )
    assert tabs is not None
    assert changed_metrics is not None

    if experiment_id == "full_experiment_id":
        assert list(p.title for p in tabs.tabs) == [
            "Summary",
            "Planner",
            "Prediction",
            "Perception",
            "Tracker",
            "VH6 2ws vs 4ws (shaA)",
            "VH6 2ws vs 4ws (shaB)",
        ], "Did not match expected tab outputs"
        assert changed_metrics == {
            "Perception": 0,
            "Planner": 0,
            "Prediction": 0,
            "Tracker": 0,
            "VH6 2ws vs 4ws": 1,
        }, "Did not match expected changed metrics"
    elif experiment_id == "steering_experiment_id":
        assert list(p.title for p in tabs.tabs) == [
            "Summary",
            "VH6 2ws vs 4ws (shaA)",
            "VH6 2ws vs 4ws (shaB)",
        ], "Did not match expected tab outputs"
        assert changed_metrics == {
            "VH6 2ws vs 4ws": 1
        }, "Did not match expected changed metrics"
    elif experiment_id == "experiment_id":
        assert list(p.title for p in tabs.tabs) == [
            "Summary",
            "Planner",
            "Prediction",
            "Perception",
            "Tracker",
        ], "Did not match expected tab outputs"
        assert changed_metrics == {
            "Perception": 0,
            "Planner": 0,
            "Prediction": 0,
            "Tracker": 0,
        }, "Did not match expected changed metrics"
    else:
        assert False, "Could not check expected tab outputs"

    assert isinstance(tabs, Tabs)


if __name__ == "__main__":
    args = [os.path.dirname(__file__), "--color=yes", "--verbose", "--verbose"]
    code = pytest.main(args)
    sys.exit(code)