import data_stat
from data_stat import Stat, Cpu

import argparse
import os, sys
import subprocess, signal
import json
import time

parser = argparse.ArgumentParser()
parser.add_argument('cpu', choices=[x.name.lower() for x in Cpu])
parser.add_argument('matrix_file')
parser.add_argument('iterations', type=int)
parser.add_argument('baseline_time_s', type=int)
parser.add_argument('baseline_delay_s', type=int)
parser.add_argument('--perf', action='store_true')
parser.add_argument('--power', action='store_true')
parser.add_argument('-d', '--debug', action='store_true')
args = parser.parse_args()
args.cpu = Cpu[args.cpu.upper()]

python = {
        Cpu.ALTRA: 'python3',
        Cpu.EPYC_7313P: 'python3.11'
}
program = {
        Cpu.ALTRA: [
            'apptainer', 'run', 'pytorch-altra.sif', '-c',
            'numactl --cpunodebind=0 --membind=0 '
            + f'python spmv.py {args.matrix_file} {args.iterations}'],
        Cpu.EPYC_7313P: [
            'apptainer', 'run', 'pytorch-epyc_7313p.sif',
            'python3', 'spmv.py', f'{args.matrix_file}', f'{args.iterations}']
}
perf = ['perf', 'stat']
perf_args = {
        Cpu.ALTRA: [
            ['-d', '-d'],
            ['-M', 'branch_misprediction_ratio'],
            ['-M', 'dtlb_walk_ratio,itlb_walk_ratio'],
            ['-M', 'l1d_cache_miss_ratio,l1i_cache_miss_ratio'],
            ['-M', 'l2_cache_miss_ratio,l2_tlb_miss_ratio,ll_cache_read_miss_ratio']]
}

def baseline_power(cpu: Cpu, baseline_time_s: int) -> list:
    power_process = subprocess.Popen([python[args.cpu], 'power.py', '-s', str(baseline_time_s)],
            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
    return [float(x) for x in power_process.communicate()[0].strip().split('\n') if len(x) != 0]

def run_program(program: list[str]) -> tuple[dict, str]:
    if args.debug:
        print(program)
    process = subprocess.run(program,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if args.debug:
        print(process.stdout)
        print(process.stderr)
    return (json.loads(process.stdout), process.stderr)

result = dict()
result[Stat.CPU.name] = args.cpu.name
result[Stat.ITERATIONS.name] = args.iterations

program_result = run_program(program[args.cpu])
result |= program_result[0]
print(program_result[1], file=sys.stderr)

if args.power:
    result[Stat.BASELINE_TIME_S.name] = args.baseline_time_s
    result[Stat.BASELINE_DELAY_S.name] = args.baseline_delay_s

    time.sleep(args.baseline_delay_s)
    result[Stat.POWER_BEFORE.name] = baseline_power(args.cpu, args.baseline_time_s)
    if args.debug:
        print(result)

    run_program(program[args.cpu]) # Warmup

    if args.cpu == Cpu.ALTRA:
        power_process = subprocess.Popen(
                [python[args.cpu], 'power.py'],
                stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)

        run_program(program[args.cpu])[0]

        power_process.send_signal(signal.SIGINT)
        if args.debug:
            print(result)

        result[Stat.POWER.name] = [float(x) for x in power_process.communicate()[0].strip().split('\n')]
        # Riemann Sum across the last (s) power recordings.
        from math import ceil
        result[Stat.JOULES.name] = (
                sum(result[Stat.POWER.name][-ceil(result[Stat.TIME_S.name]):-1])
                + (result[Stat.POWER.name][-1] * (result[Stat.TIME_S.name] % 1)))

    elif args.cpu == Cpu.EPYC_7313P:
        power_process = subprocess.Popen(
                [python[args.cpu], 'power.py'] + program[args.cpu],
                stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
        result[Stat.POWER.name] = [float(x) for x in power_process.communicate()[0].strip().split('\n')]
        result[Stat.JOULES.name] = result[Stat.POWER.name][0] * result[Stat.TIME_S.name]

    if args.debug:
        print(result)
        #print(len(result['power']))
        #print(sum(result['power']) / len(result['power']))

    time.sleep(args.baseline_delay_s)
    result[Stat.POWER_AFTER.name] = baseline_power(args.cpu, args.baseline_time_s)
    if args.debug:
        print(result)

if args.perf:
    for perf_arg in perf_args[args.cpu]:
        output = run_program(perf + perf_arg + program[args.cpu])[1]
        print(output, file=sys.stderr)
        result = result | data_stat.parse_output(output, args.cpu)
        if args.debug:
            print(result)

    result = result | data_stat.derive_stats(result)

    if args.debug:
        print(result)

print(json.dumps(result))

#arch = subprocess.run(['uname', '-m'], stdout=subprocess.PIPE, text=True).stdout.strip()
#baseline = subprocess.run(
#        ['./power.sh', args.baseline_time_s],
#        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
#print(baseline)
#for line in baseline.stdout.split('\n'):
#    print("line")
#    print(line)

#os.path.basename(args.matrix_file)

#subprocess.run(
#    ['apptainer', 'run', 'pytorch-altra.sif', '-c',
#        f'"numactl --cpunodebind=0 --membind=0 python spmv.py {args.matrix_file} {args.iterations}"'
#    ],
#    stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)