#!/usr/bin/env python3

from .main import *

def _test_schedule(generated_schedule, literal_schedule):
    """Helper function for `test_schedules`."""
    diff = set(generated_schedule.keys()) - set(literal_schedule.keys())
    try:
        assert len(diff) == 0
    except:
        print(f"Generated schedule contains an impossible quantity of assemblers: {diff}")

    for n in sorted(literal_schedule.keys()):
        try:
            assert generated_schedule[n][0] == literal_schedule[n][0]
            assert generated_schedule[n][1] == literal_schedule[n][1]
        except:
            print(f"Generated schedule is wrong for {n} assemblers")
            print(f"Is:        {generated_schedule[n]}")
            print(f"Should be: {literal_schedule[n]}")

def test_schedules():
    # check that only one strategy exists for schedules with too few hours
    for i in range(HOURS_TO_PRODUCE_ASSEMBLER):
        res = assembler_schedule(i)
        _test_schedule(res, { 1: (i, 0) })

    # check that complex strategies can be formed
    res = assembler_schedule(HOURS_TO_PRODUCE_ASSEMBLER)
    _test_schedule(res, {
        1: (HOURS_TO_PRODUCE_ASSEMBLER, 0),
        2: (0, 0),
    })

    # check handling of remaining hours without fighter production potential
    res = assembler_schedule(HOURS_TO_PRODUCE_ASSEMBLER + 1)
    _test_schedule(res, {
        1: (HOURS_TO_PRODUCE_ASSEMBLER + 1, 0),
        2: (1, 0),
    })

    # check handling of remaining hours with fighter production potential
    res = assembler_schedule(HOURS_TO_PRODUCE_ASSEMBLER * 2)
    _test_schedule(res, {
        1: (HOURS_TO_PRODUCE_ASSEMBLER * 2, 0),
        2: (HOURS_TO_PRODUCE_ASSEMBLER, 0),
        3: (0, FIGHTERS_PRODUCED_IN_ONE_UNIT),
        4: (0, 0),
    })

    # check that non-trivial ranges are bounded correctly
    res = assembler_schedule(HOURS_TO_PRODUCE_ASSEMBLER * 3)
    _test_schedule(res, {
        1: (HOURS_TO_PRODUCE_ASSEMBLER * 3, 0),
        2: (HOURS_TO_PRODUCE_ASSEMBLER * 2, 0),
        3: (HOURS_TO_PRODUCE_ASSEMBLER, FIGHTERS_PRODUCED_IN_ONE_UNIT),
        4: (HOURS_TO_PRODUCE_ASSEMBLER, 0),
        5: (0, FIGHTERS_PRODUCED_IN_ONE_UNIT*3),
        6: (0, FIGHTERS_PRODUCED_IN_ONE_UNIT*2),
        7: (0, FIGHTERS_PRODUCED_IN_ONE_UNIT),
        8: (0, 0),
    })

def main():
    print("Running test suite...")
    test_schedules()

if __name__ == "__main__":
    main()