#!/usr/bin/env python3

import argparse
import json
import os
import statistics
from collections import defaultdict
from tools.stats.s3_stat_parser import (
    get_previous_reports_for_branch,
    Report,
    Version2Report,
)
from typing import cast, DefaultDict, Dict, List, Any
from urllib.request import urlopen

SLOW_TESTS_FILE = ".pytorch-slow-tests.json"
SLOW_TEST_CASE_THRESHOLD_SEC = 60.0
RELATIVE_DIFFERENCE_THRESHOLD = 0.1
IGNORED_JOBS = ["asan", "periodic"]


def get_test_case_times() -> Dict[str, float]:
    reports: List[Report] = get_previous_reports_for_branch("origin/viable/strict", "")
    # an entry will be like ("test_doc_examples (__main__.TestTypeHints)" -> [values]))
    test_names_to_times: DefaultDict[str, List[float]] = defaultdict(list)
    for report in reports:
        if report.get("format_version", 1) != 2:  # type: ignore[misc]
            raise RuntimeError("S3 format currently handled is version 2 only")
        v2report = cast(Version2Report, report)

        if any(job_name in str(report["build_job"]) for job_name in IGNORED_JOBS):
            continue

        for test_file in v2report["files"].values():
            for suitename, test_suite in test_file["suites"].items():
                for casename, test_case in test_suite["cases"].items():
                    # The below attaches a __main__ as that matches the format of test.__class__ in
                    # common_utils.py (where this data will be used), and also matches what the output
                    # of a running test would look like.
                    name = f"{casename} (__main__.{suitename})"
                    succeeded: bool = test_case["status"] is None
                    if succeeded:
                        test_names_to_times[name].append(test_case["seconds"])
    return {
        test_case: statistics.mean(times)
        for test_case, times in test_names_to_times.items()
    }


def filter_slow_tests(test_cases_dict: Dict[str, float]) -> Dict[str, float]:
    return {
        test_case: time
        for test_case, time in test_cases_dict.items()
        if time >= SLOW_TEST_CASE_THRESHOLD_SEC
    }


def get_test_infra_slow_tests() -> Dict[str, float]:
    url = "https://raw.githubusercontent.com/pytorch/test-infra/generated-stats/stats/slow-tests.json"
    contents = urlopen(url, timeout=1).read().decode("utf-8")
    return cast(Dict[str, float], json.loads(contents))


def too_similar(
    calculated_times: Dict[str, float], other_times: Dict[str, float], threshold: float
) -> bool:
    # check that their keys are the same
    if calculated_times.keys() != other_times.keys():
        return False

    for test_case, test_time in calculated_times.items():
        other_test_time = other_times[test_case]
        relative_difference = abs(
            (other_test_time - test_time) / max(other_test_time, test_time)
        )
        if relative_difference > threshold:
            return False
    return True


def export_slow_tests(options: Any) -> None:
    filename = options.filename
    if os.path.exists(filename):
        print(f"Overwriting existent file: {filename}")
    with open(filename, "w+") as file:
        slow_test_times: Dict[str, float] = filter_slow_tests(get_test_case_times())
        if options.ignore_small_diffs:
            test_infra_slow_tests_dict = get_test_infra_slow_tests()
            if too_similar(
                slow_test_times, test_infra_slow_tests_dict, options.ignore_small_diffs
            ):
                slow_test_times = test_infra_slow_tests_dict
        json.dump(
            slow_test_times, file, indent="    ", separators=(",", ": "), sort_keys=True
        )
        file.write("\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Export a JSON of slow test cases in PyTorch unit test suite"
    )
    parser.add_argument(
        "-f",
        "--filename",
        nargs="?",
        type=str,
        default=SLOW_TESTS_FILE,
        const=SLOW_TESTS_FILE,
        help="Specify a file path to dump slow test times from previous S3 stats. Default file path: .pytorch-slow-tests.json",
    )
    parser.add_argument(
        "--ignore-small-diffs",
        nargs="?",
        type=float,
        const=RELATIVE_DIFFERENCE_THRESHOLD,
        help="Compares generated results with stats/slow-tests.json in pytorch/test-infra. If the relative differences "
        "between test times for each test are smaller than the threshold and the set of test cases have not "
        "changed, we will export the stats already in stats/slow-tests.json. Else, we will export the calculated "
        "results. The default threshold is 10%.",
    )
    return parser.parse_args()


def main() -> None:
    options = parse_args()
    export_slow_tests(options)


if __name__ == "__main__":
    main()

[(23, 'tools.stats.s3_stat_parser.get_previous_reports_for_branch', 'get_previous_reports_for_branch', False, 'from tools.stats.s3_stat_parser import get_previous_reports_for_branch, Report, Version2Report\n'), (25, 'collections.defaultdict', 'defaultdict', False, 'from collections import defaultdict\n'), (83, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (100, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (29, 'typing.cast', 'cast', False, 'from typing import cast, DefaultDict, Dict, List, Any\n'), (45, 'statistics.mean', 'statistics.mean', False, 'import statistics\n'), (61, 'json.loads', 'json.loads', False, 'import json\n'), (93, 'json.dump', 'json.dump', False, 'import json\n'), (60, 'urllib.request.urlopen', 'urlopen', False, 'from urllib.request import urlopen\n')]