#!/usr/bin/env python

"""VisualizeTestTimes.py

Visualize CTest test times with VTK.

Run from the top of the build tree after the ctest has been run at least once.

Pass the --modular-grouping flag to group by module."""


import os
import pprint
import subprocess
import sys
import vtk

vtk_major_version = vtk.vtkVersion.GetVTKMajorVersion()

if len(sys.argv) > 1 and sys.argv[1] == "-h":
    print("Usage: VisualizeTestTimes.py [--modular-grouping]")
    sys.exit(1)
modular = False
if len(sys.argv) > 1 and sys.argv[1] == "--modular-grouping":
    modular = True

# Sanity check to ensure we are in the build tree
test_cost_data_file = os.path.join("Testing", "Temporary", "CTestCostData.txt")
if not os.path.exists(test_cost_data_file):
    print("Run this script from the build tree after running ctest " + "at least once.")
    sys.exit(1)

# Read the input data
with open(test_cost_data_file) as fp:
    test_cost_data_lines = fp.readlines()
failed_tests_index = test_cost_data_lines.index("---\n")

# Import the data into a vtkTable
table = vtk.vtkTable()
id_array = vtk.vtkUnsignedIntArray()
id_array.SetName("Pedigree Id")
table.AddColumn(id_array)
attributes = table.GetAttributes(vtk.vtkDataObject.ROW)
attributes.SetActivePedigreeIds("Pedigree Id")
test_name_array = vtk.vtkStringArray()
test_name_array.SetName("Test Name")
table.AddColumn(test_name_array)
number_of_runs_array = vtk.vtkUnsignedIntArray()
number_of_runs_array.SetName("Number of Runs")
table.AddColumn(number_of_runs_array)
test_time_array = vtk.vtkFloatArray()
test_time_array.SetName("Test Time")
table.AddColumn(test_time_array)
runs_long_array = vtk.vtkStringArray()
runs_long_array.SetName("RUNS_LONG Label")
table.AddColumn(runs_long_array)
other_labels_array = vtk.vtkStringArray()
other_labels_array.SetName("Other Labels")
table.AddColumn(other_labels_array)
table.SetNumberOfRows(failed_tests_index)
ctest_exe = "ctest"
runs_long = subprocess.check_output([ctest_exe, "-L", "RUNS_LONG", "-N"])
runs_long = runs_long.split("\n")[1:-3]
runs_long = [ii.split()[2] for ii in runs_long]
has_runs_long_time = 0.0
no_runs_long_time = 0.0
for ii in range(failed_tests_index):
    split = test_cost_data_lines[ii].strip().split()
    table.SetValue(ii, 0, ii)
    name = split[0]
    table.SetValue(ii, 1, name)
    table.SetValue(ii, 2, int(split[1]))
    time = float(split[2])
    table.SetValue(ii, 3, time)
    if name in runs_long:
        table.SetValue(ii, 4, "Has RUNS_LONG Label")
        has_runs_long_time += time
    else:
        table.SetValue(ii, 4, "No RUNS_LONG Label")
        no_runs_long_time += time
    table.SetValue(ii, 5, "None")
labels = subprocess.check_output([ctest_exe, "--print-labels"])
labels = labels.split("\n")[2:-1]
labels = [ii.strip() for ii in labels]
if "RUNS_LONG" in labels:
    labels.pop(labels.index("RUNS_LONG"))
# Assuming tests will only have RUNS_LONG and up to only one other label
if modular:
    for label in labels:
        tests = subprocess.check_output([ctest_exe, "-L", label, "-N"])
        tests = tests.split("\n")[2:-3]
        tests = [ii.split()[2] for ii in tests]
        for test in tests:
            index = test_name_array.LookupValue(test)
            other_labels_array.SetValue(index, label)
print("RUNS_LONG tests:")
pprint.pprint(runs_long)
print(
    "RUNS_LONG time percentage: {:.4}%".format(
        str(has_runs_long_time / (has_runs_long_time + no_runs_long_time) * 100)
    )
)
print(
    "RUNS_LONG test percentage: {:.4}%".format(
        str((len(runs_long)) / (failed_tests_index - 1.0) * 100)
    )
)

# Convert the vtkTable to a vtkTree
table_to_tree = vtk.vtkTableToTreeFilter()
if vtk_major_version == 5:
    table_to_tree.SetInput(table)
else:
    table_to_tree.SetInputData(table)
group_runs_long = vtk.vtkGroupLeafVertices()
group_runs_long.SetInputConnection(table_to_tree.GetOutputPort())
group_runs_long.SetInputArrayToProcess(
    0, 0, 0, vtk.vtkDataObject.VERTEX, "RUNS_LONG Label"
)
group_runs_long.SetInputArrayToProcess(1, 0, 0, vtk.vtkDataObject.VERTEX, "Test Name")
group_other_label = vtk.vtkGroupLeafVertices()
group_other_label.SetInputConnection(group_runs_long.GetOutputPort())
group_other_label.SetInputArrayToProcess(
    0, 0, 0, vtk.vtkDataObject.VERTEX, "Other Labels"
)
group_other_label.SetInputArrayToProcess(1, 0, 0, vtk.vtkDataObject.VERTEX, "Test Name")

# Visualize with a tree map view
tree_map_view = vtk.vtkTreeMapView()
if modular:
    tree_map_view.SetTreeFromInputConnection(group_other_label.GetOutputPort())
else:
    tree_map_view.SetTreeFromInputConnection(group_runs_long.GetOutputPort())
tree_map_view.SetAreaLabelArrayName("Test Name")
tree_map_view.SetAreaHoverArrayName("Test Name")
tree_map_view.SetAreaLabelVisibility(True)
tree_map_view.SetAreaSizeArrayName("Test Time")
tree_map_view.DisplayHoverTextOn()
tree_map_view.SetLayoutStrategyToSquarify()

# Pretty preference: Mellow, Neon, Ocean
theme = vtk.vtkViewTheme.CreateNeonTheme()
tree_map_view.ApplyViewTheme(theme)

tree_map_view.Update()
tree_map_view.ResetCamera()

interactor = tree_map_view.GetInteractor()
interactor.Initialize()
interactor.Start()