import argparse
import zipfile
import json
import os

# Costs
CONSTANT_COST = 1
REGISTER_COST = 1
OP_COST = {
    "SUB": 4,
    "ADD": 4,
    "INC": 4,
    "DEC": 4,
    "MUL": 10,
    "DIV": 10,
    "MOD": 10,
    "OR": 8,
    "AND": 8,
    "XOR": 8,
    "INV": 8,
    "SL": 5,
    "SR": 5,
    "SLU": 5,
    "SRU": 5,
    "ROTL": 7,
    "ROTR": 7,
}
ALG_LINE_COST = 0.5

# ansi escape codes
ul = "\033[4m"  # underline
end = "\033[0m"  # reset
ylw = "\033[33m"  # yellow


def is_empty_row(row, pedantic=False):
    # Check if signal table is non-empty
    if "signal" in row and len(row["signal"]) != 0:
        return False

    # Only check if keys are set if pedantic is true
    if not pedantic:
        return True
    # Check if "unconditional-jump", "conditional-jump" or "label" is set
    keys = ("unconditional-jump", "conditional-jump", "label")
    return not any(key in row for key in keys)


def evaluate(filepath, verbose=False, pedantic=False):
    filename = filepath.split(os.path.sep)[-1]
    if not filepath.endswith(".zip"):
        print(
            f"{filename} :: Supplied file does not have .zip file extension. Skipping .."
        )
        return -1

    with zipfile.ZipFile(filepath, "r") as savefile:
        with savefile.open("machine.json", "r") as machinefile, savefile.open(
            "signal.json", "r"
        ) as signalfile:
            machine = json.load(machinefile)
            signal = json.load(signalfile)

    # Load lines of code (rows)
    total_rows = signal["signaltable"]["row"]
    # Remove rows without effect (used for formatting for example)
    rows = [row for row in total_rows if not is_empty_row(row, pedantic=pedantic)]
    if verbose:
        print(f"{filename} :: Total number of rows: {len(total_rows)}")
        print(f"{filename} :: Number of rows after excluding empty: {len(rows)}")

    # Check if IR or PC register was used

    pc_used = any(
        (
            signal["name"] == "PC.W" and signal["value"] == "1"
            for row in rows
            for signal in row["signal"]
        )
    )
    ir_used = any(
        (
            signal["name"] == "IR.W" and signal["value"] == "1"
            for row in rows
            for signal in row["signal"]
        )
    )

    if verbose:
        if pc_used:
            print(f"{filename} :: PC Register was used in signal table row.")
        if ir_used:
            print(f"{filename} :: IR Register was used in signal table row.")

    # Load used multiplexer constants
    try:
        mux_input_a = next(
            filter(lambda mux: mux["muxType"] == "A", machine["machine"]["muxInputs"])
        )["input"]
        mux_consts_a = [
            int(mux_input["value"])
            for mux_input in mux_input_a
            if mux_input["type"] == "constant"
        ]
    except StopIteration:
        print(
            f"{filename} :: Couldn't find input for multiplexer A. Is the file corrupted? Skipping file .."
        )
        return -1
    try:
        mux_input_b = next(
            filter(lambda mux: mux["muxType"] == "B", machine["machine"]["muxInputs"])
        )["input"]
        mux_consts_b = [
            int(mux_input["value"])
            for mux_input in mux_input_b
            if mux_input["type"] == "constant"
        ]
    except StopIteration:
        print(
            f"{filename} :: Couldn't find input for multiplexer B. Is the file corrupted? Skipping file .."
        )
        return -1

    # Base machine has constants 0 and 1 at multiplexer A. All other constants are extensions.
    base_muxt_a = (0, 1)
    for base_input in base_muxt_a:
        try:
            mux_consts_a.remove(base_input)
        except ValueError:
            pass

    constants = set(mux_consts_a + mux_consts_b)

    if verbose:
        print(
            f"{filename} :: Found {len(mux_consts_a)} constants for multiplexer A: {mux_consts_a}"
        )
        print(
            f"{filename} :: Found {len(mux_consts_b)} constants for multiplexer B: {mux_consts_b}"
        )
        print(
            f"{filename} :: Found {len(constants)} total unique constants: [{', '.join([str(c) for c in constants])}]"
        )

    # Load used registers
    registers = machine["machine"]["registers"]["register"]
    registers = [register["name"] for register in registers]

    if pc_used:
        registers.append("PC_ALT")
    if ir_used:
        registers.append("IR_ALT")

    if verbose:
        print(f"{filename} :: Found {len(registers)} additional registers: {registers}")

    # Load used operations
    operations = machine["machine"]["alu"]["operation"]
    base_operations = ("A_ADD_B", "B_SUB_A", "TRANS_A", "TRANS_B")
    for base_op in base_operations:
        try:
            operations.remove(base_op)
        except ValueError:
            pass

    # Extract operation (remove operands)
    def get_op(op_str):
        return next(filter(lambda substr: substr not in ("A", "B"), op_str.split("_")))

    operations = list(map(get_op, operations))

    if verbose:
        print(
            f"{filename} :: Found {len(operations)} additional operations: {operations}"
        )

    # Sum points
    alg_line_costs = len(rows) * ALG_LINE_COST  # every line of code
    constant_costs = len(constants) * CONSTANT_COST  # constants at both multiplexers
    register_costs = len(registers) * REGISTER_COST  # registers
    operation_costs = 0
    for operation in operations:  # operations
        operation_costs += OP_COST[operation]

    total = alg_line_costs + constant_costs + register_costs + operation_costs

    # Summarize
    costs = (alg_line_costs, constant_costs, register_costs, operation_costs, total)
    precision = max(
        [len(str(float(cost)).split(".")[1].lstrip("0")) for cost in costs]
    )  # unreadable but works ¯\_(ツ)_/¯

    if verbose:
        print("")

    print(f"{ul}Summary for {filename}:{end}\n")
    print(f"   {alg_line_costs:5.{min(precision, 2)}f} LINES (excluding empty lines)")
    print(f"+  {constant_costs:5.{min(precision, 2)}f} CONSTANTS")
    print(f"+  {register_costs:5.{min(precision, 2)}f} REGISTERS")
    print(f"+  {operation_costs:5.{min(precision, 2)}f} OPERATIONS")
    print(f"-------------")
    print(f"=  {ylw}{total:5.{min(precision, 2)}f} TOTAL{end}\n\n")

    return total


# Evaluation
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "source",
        type=str,
        nargs="+",
        help="Either ZIP file(s) generated by simulator or the submission root folder",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        dest="verbose",
        action="store_true",
        help="Prints additional information.",
    )
    parser.add_argument(
        "-p",
        "--pedantic",
        dest="pedantic",
        action="store_true",
        help="Extra pedantic (for example when checking for empty lines)",
    )
    parser.add_argument(
        "-t",
        "--top",
        dest="top",
        type=int,
        help="Print top n candidates (defaults to 7)",
    )
    args = parser.parse_args()

    verbose = args.verbose
    pedantic = args.pedantic
    top_n = args.top if args.top != None else 7

    # output score for each file
    scores = []

    # check if source argument is folder
    savefiles = []
    for source in args.source:
        if os.path.isdir(source):
            # add all zip from subdirectories
            for d in [
                e for e in os.listdir(source) if os.path.isdir(os.path.join(source, e))
            ]:
                savefiles.append(os.path.join(source, d, f"{d}.zip"))
        elif source.endswith(".zip"):
            savefiles.append(source)
        else:
            print(f"Source '{source}' is not a ZIP file.")

    for savefile in savefiles:
        score = evaluate(savefile, verbose=verbose, pedantic=pedantic)
        if score == -1:
            continue

        scores.append([savefile, score])

    # if there is more than one file, output top 3
    scores.sort(key=lambda x: x[1])

    n = len(savefiles)
    if n > 1:
        print(f"{ul}Leaderboard:{end}")
        for i in range(min(n, top_n)):
            file, score = scores[i]
            print(f"#{i + 1} - {ylw}{score:5.2f} TOTAL{end} - {file}")