#!/usr/bin/python3
import sys
import pythia
import random
import json
import os

from subprocess import DEVNULL, check_call

def load_coverage(fn):
    with open(fn, "r") as f:
        lines = [l.strip() for l in f.readlines()] 
        lines = {i: c for (i, c) in enumerate(lines)}
        return lines

def reduce_tests(test_coverage, k=5):
    selection = set()
    line_coverage = {i: 0 for (i, v) in enumerate(cov[0]) if v != 'X'}

    for i in range(1, k + 1):
        for l1 in line_coverage:
            if line_coverage[l1] >= i:
                continue
            
            # find the sub-set of tests that cover this line
            candidates = set(t for (t, tcov) in test_coverage.items() if tcov[l1] == '1')

            # if nothing covers this line, skip!
            if not candidates:
                continue

            # find which lines don't meet the desired coverage level for this iteration
            poorly_covered = set(l for (l, c) in line_coverage.items() if c < i)

            # find the sub-set of candidates that cover the greatest number of these
            # underrepresented lines
            best_candidates = set()
            max_coverage = 0
            for t in candidates:
                t_coverage = sum(1 for l2 in poorly_covered if test_coverage[t][l2] == '1')
                if t_coverage > max_coverage:
                    max_coverage = t_coverage
                    best_candidates = set([t])
                elif t_coverage == max_coverage:
                    best_candidates.add(t)

            # randomly pick a test from the best candidates
            test = random.choice(tuple(best_candidates))
            tcov = test_coverage[test]

            # update the line coverage
            for l2 in line_coverage:
                if tcov[l2] == '1':
                    line_coverage[l2] += 1

            # remove this test from the pool
            selection.add(test)
            del test_coverage[test]

    return selection


def reduce_tests_random(test_coverage, k=2):
    selection = set()
    line_coverage = {i: 0 for (i, v) in enumerate(cov[0]) if v != 'X'}

    for line in line_coverage:
        if line_coverage[line] >= k:
            continue

        # select a sub-set of tests that cover this line at random, such that
        # at least k tests cover this line
        candidates = set(t for (t, tcov) in test_coverage.items() if tcov[line] == '1')
        if candidates == set():
            print("WARNING: failed to cover line {}".format(line))
            continue
        sample_size = min(len(candidates), k - line_coverage[line])
        sample = random.sample(candidates, sample_size)

        # remove these tests from further consideration, then add to selection
        for t in sample:
            del test_coverage[t]
            selection.add(t)

        # update line coverage
        line_coverage[line] += sample_size

    return selection


def reduce_tests_by_patterns(test_coverage, k=1):
    # group tests by coverage patterns
    trace_cov = {}
    for (test, tcov) in test_coverage.items():
        if not tcov in trace_cov:
            trace_cov[tcov] = set()
        trace_cov[tcov].add(test)
        
    # for each coverage pattern, select k tests
    selection = []
    for tests in trace_cov.values():
        tests = random.sample(tests, min(len(tests), k))
        selection += tests
    selection = set(selection)

    return selection


"""
Reduces a given test suite, based on provided coverage, such that each line,
where possible, is covered by +k+ tests.
"""
def basic_reduce_tests(test_coverage, k=1):
    selection = set()
    line_coverage = {i: 0 for (i, v) in enumerate(cov[0]) if v != 'X'}

    for line_no in line_coverage:
        if line_coverage[line_no] >= k: 
            continue
        
        # find the sub-set of remaining tests that cover this line
        candidates = set(t for (t, tcov) in test_coverage.items() if tcov[line_no] == '1')
        if candidates == set():
            print("WARNING: failed to cover line {}".format(line_no))
            continue

        # find the (set of) tests that cover the most uncovered lines
        best_candidates = set()
        max_coverage = 0
        for t in candidates:
            t_coverage = test_coverage[t][line_no:].count('1')
            if t_coverage > max_coverage:
                max_coverage = t_coverage
                best_candidates = set([t])
            elif t_coverage == max_coverage:
                best_candidates.add(t)

        # randomly pick a test from the best candidates
        test = random.choice(tuple(best_candidates))
        tcov = test_coverage[test]

        # update the line coverage
        for i in line_coverage:
            if tcov[i] == '1':
                line_coverage[i] += 1

        # remove this test from the pool
        selection.add(test)
        del test_coverage[test]

    return selection


"""
Rebuilds the test manifest and oracle for this problem, using a subset of the
original test suite.
"""
def rebuild_manifest_and_oracle(tests):

    # find the tests in the original manifest
    manifest = pythia.TestManifest("tests.pythia.json")
    original_tests = manifest.contents()
    tests = [original_tests[t] for t in tests]
    tests = [t.to_json() for t in tests]
    
    # save the original files to the oracle directory
    os.mkdir("oracle")
    os.rename("tests.pythia.json", "oracle/tests.pythia.json")
    os.rename("oracle.pythia.json", "oracle/oracle.pythia.json")

    # write the new test manifest to file
    with open("tests.pythia.json", "w") as f:
        json.dump(tests, f)

    # re-generate the oracle
    exe = os.environ['EXECUTABLE']
    check_call("./compile.sh source/{}.".format(exe), stdout=DEVNULL, stderr=DEVNULL, shell=True)
    os.rename("source/{}".format(exe), exe)
    pythia.generate(exe)


if __name__ == "__main__":
    random.seed(0)
    cov = load_coverage("coverage.txt")
    k = int(sys.argv[1])
    tests = reduce_tests(cov, k)
    rebuild_manifest_and_oracle(tests)
