#!/usr/bin/python3
#####
# Copyright 2020 Joshua Bakita
#
# This program generates input data for the DIS benchmark suite on stdout
# given a requested working set size.
#####


from ctypes import sizeof, c_double, c_int, c_short
from math import sqrt, floor
import sys # For argv and stderr

USAGE = """Usage: {} <benchmark> <template> <WSS in bytes>"""

# Check input
if (len(sys.argv) < 4):
    print(USAGE.format(sys.argv[0]), file=sys.stderr)
    exit(1);

# Don't try to understand the logic in these functions, see WSS_DOCS.md
def setup_field(params, wss):
    params[0] = wss
    return params

def setup_matrix(params, wss):
    nnZR = 0.08 # 8% seems average
    # This formula is out of a solver
    si = sizeof(c_int)
    sd = sizeof(c_double)
    d = (sqrt((si**2) * (-(nnZR-1)) - si*sd*(nnZR-9) + si*wss*nnZR + sd*(25*sd+wss*nnZR+wss)) - si - 5*sd) / (si*nnZR + sd*nnZR + sd)
    params[1] = floor(d);
    params[2] = floor(d*d*nnZR);
    if params[1] <= 0 or params[2] <= 0:
        raise Exception("WSS too small for matrix benchmark!")
    return params

def setup_neighborhood(params, wss):
    bitDepth = 8
    bitDepthAlloc = sizeof(c_int) * (2**(bitDepth + 1) - 1)
    dim = sqrt((wss - bitDepthAlloc) / sizeof(c_short))
    params[1] = bitDepth
    params[2] = floor(dim)
    if params[1] <= 0 or params[2] <= 0:
        raise Exception("WSS too small for neighborhood benchmark!")
    # Cap maximum line thinkness to the image size
    params[5] = min(params[2]-1, int(params[5]))
    # Cap line lengths to the image size
    params[6] = min(params[2]-1, int(params[6]))
    params[7] = min(params[2]-1, int(params[7]))
    return params

def setup_pointer(params, wss):
    n = 10;
    f = (wss - sizeof(c_int) * 4 * n) / sizeof(c_int)
    params[0] = floor(f)
    params[4] = floor(n)
    if params[0] <= 0 or params[4] <= 0:
        raise Exception("WSS too small for pointer benchmark!")
    return params

def setup_transitive(params, wss):
    n = sqrt(wss / (sizeof(c_int) * 2))
    params[0] = floor(n)
    # Fix edges at 50%
    params[1] = floor(params[0] * 0.5)
    if params[0] <= 0:
        raise Exception("WSS too small for transitive benchmark!")
    return params

def setup_update(params, wss):
    f = wss / sizeof(c_int)
    params[0] = floor(f)
    if params[0] <= 0:
        raise Exception("WSS too small for update benchmark!")
    # Don't do more than 100M hops (keeps time array feasible)
    params[2] = min(100000000, int(params[2]))
    # Enforce size requirements
    params[4] = min(params[0]-1, int(params[4]))
    params[5] = min(params[0]-1, int(params[5]))
    params[6] = min(params[0]-1, int(params[6]))
    return params

def setup_random_walk(params, wss):
    params[0] = wss
    return params

BENCH_TO_PARAMS = {"field":setup_field, "matrix":setup_matrix, "neighborhood":setup_neighborhood, "pointer":setup_pointer, "transitive":setup_transitive, "update":setup_update, "random_walk":setup_random_walk}

# Main logic
benchmark_name = sys.argv[1]
if benchmark_name not in BENCH_TO_PARAMS.keys():
    print("Invalid benchmark name.", file=sys.stderr)
    exit(2)

wss = int(sys.argv[3])
if wss <= 0:
    print("Invalid working set size", file=sys.stderr)
    exit(3)

with open(sys.argv[2], "r") as template:
    # We expect the initialization params to all be on the first line
    params = template.readline().split()
    mutated_params = BENCH_TO_PARAMS[benchmark_name](params, wss);
    print(" ".join(map(lambda x: str(x), mutated_params)))
    print(" ".join(map(lambda x: str(x), mutated_params)), file=sys.stderr)
    if benchmark_name == "pointer":
        # Clone the data format used in the template
        for i in range(0,10):
            print("10 " + str(mutated_params[0]-10) + " " + str(mutated_params[0]-10))
    else:
        print(template.read())