import re
from shpc_grader import SHPCProblem, SHPCTestCase
def correct_check(result):
    result = [line for line in result.split('\n') if re.search(r"Result", line)][0]
    result = result.split(' ')[1]
    result = result.strip()
    return result == "VALID"

def correct_check_currying(score):
    return lambda result: score if correct_check(result) else 0

def correct_and_perf_check(result):
    if correct_check(result):
        result = [line for line in result.split('\n') if re.search(r"throughput", line)][0]
        result = result.split(' ')[2]
        result = result.strip()
        return float(result)
    return 0

def correct_and_perf_check_currying(maxperf, score):
    return lambda result: min(maxperf, correct_and_perf_check(result)) * score / maxperf

hw2_2 = SHPCProblem("hw2-2", "hw2", "hw2/vectordot", 
      [ "vectordot.c" ],
      [ 
          SHPCTestCase("./main -m fma -n 1 1", correct_check_currying(4)),
          SHPCTestCase("./main -m fma -n 1 3", correct_check_currying(4)),
          SHPCTestCase("./main -m fma -n 1 11", correct_check_currying(4)),
          SHPCTestCase("./main -m fma -n 1 4097", correct_check_currying(4)),
          SHPCTestCase("./main -m fma -n 1 1000000", correct_check_currying(4)),
      ],
    )

hw2_3 = SHPCProblem("hw2-3", "hw2", "hw2/matmul", 
      [ "matmul.c" ],
      [ 
          # Accuracy test
          SHPCTestCase("./main -v -t 26 831 538 2304", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 9 3305 1864 3494", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 38 618 3102 1695", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 30 1876 3453 3590", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 16 1228 2266 1552", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 2 3347 171 688", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 39 3583 962 765", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 30 2962 373 1957", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 9 3646 2740 3053", correct_check_currying(1)),
          SHPCTestCase("./main -v -t 26 1949 3317 3868", correct_check_currying(1)),
          # Performance test
          SHPCTestCase("./main -v -t 64 4096 4096 4096", correct_and_perf_check_currying(60, 20)),
      ],
    )