import multiprocessing
import struct

from utils import measure, read_binary_file


@measure
def run_concurrent(num_files):
    files = list(range(num_files))
    with multiprocessing.Pool(processes=4) as pool:
        ret = pool.map(compute_subdomain, files)
    result = compute_result(ret)
    print("PARTIAL RESULTS:", ret)
    print("FINAL RESULTS:", result)
    return compute_result(ret)

def compute_subdomain(i):
    # 1. READ files with format [int]_input.bin
    data = read_data_i(i)
    # 2. COMPUTE
    a = data[::2]
    b = data[1::2]
    # 3. COMPUTE...
    r = compute_saxpy(a, b)
    ret = compute_result(r)
    # 4. WRITE RESULTS:
    write_data(r)
    return ret

def write_data(r):
    #waste_time_rand()
    with open("results.bin", 'wb') as f:
        for num in r:
            f.write(struct.pack('f', num))

def read_data(num_files):
    ret = []
    for i in range(num_files):
        fname = f"{i}_input.bin"
        ret.extend(read_binary_file(fname))
    return ret

def read_data_i(i: int):
    aux = f"{i}_input.bin"
    return read_binary_file(aux)

def compute_result(r):
    return sum(r)

def compute_saxpy(a, b):
    ret = []
    for i, j in zip(a, b):
        ret.append((33 * i) + j)
    return ret
