#!/usr/bin/python3

import os
import re
import subprocess
import sys
from run_experiments import run_exp
from epc_probe import get_max_epc_bytes

def get_free_gib():
    """Get the number of free GiB of memory"""
    ret = subprocess.run(["free", "-m"], capture_output=True)
    if ret.returncode != 0:
        print("Unable to run free", file=sys.stderr)
        sys.exit(1)
    match = re.search(r'Mem:\s*\d+\s+\d+\s+(\d+)', str(ret.stdout))
    if match is None:
        print("Unable to parse output of free", file=sys.stderr)
        sys.exit(1)
    free_mem_mib = int(match.group(1))
    return free_mem_mib / 1024

# How much memory is available?
free_gib = get_free_gib()

# How much EPC is available?
epc_gib = get_max_epc_bytes() / (1<<30)


LOG_FOLDER = "Experiments/"

NUM_EPOCHS = 10
TOKEN_OUT = 1
TOKEN_IN = 1
ID_OUT = 1
ID_IN = 1
# B = message size (bytes)
B = 256
N_MAX = 1<<20
M_MAX = 72

# For 1<<20 clients and 72 servers, we need 4.92 GiB of EPC and 20 GiB
# of free RAM.
# If we don't have that, abort, unless the SHRINK_TO_MEM env var is set
# to 1, in which case perform smaller experiments that should fit in the
# available resources.
if epc_gib < 4.92 or free_gib < 20:
    shrink_to_mem = \
        os.getenv("SHRINK_TO_MEM", '0').lower() in ('true', '1', 't')
    if not shrink_to_mem:
        if epc_gib < 4.92:
            print(f"""
*** Available EPC appears to be {epc_gib:.2f} GiB; 4.92 GiB is
*** needed to run the full set of experiments.
""")
        if free_gib < 20:
            print(f"""
*** Free RAM appears to be {free_gib:.2f} GiB; 20 GiB is
*** needed to run the full set of experiments.
""")
        print("""
*** Set the env var SHRINK_TO_MEM=1 to run a smaller set of
*** experiments that should fit in the available resources.
""")
        sys.exit(1)
    # Typical EPC sizes are slightly less than a power of 2.
    if epc_gib > 3.97 and free_gib > 14:
        N_MAX = 1<<19
        M_MAX = 72
    if epc_gib > 1.97 and free_gib > 14:
        N_MAX = 1<<19
        M_MAX = 36
    elif epc_gib > 0.98 and free_gib > 7.5:
        N_MAX = 1<<18
        M_MAX = 24
    elif epc_gib > 0.9 and free_gib > 4.2:
        N_MAX = 1<<17
        M_MAX = 48
    else:
        print(f"""
*** You have too small EPC ({epc_gib:.2f} GiB) and/or free
*** memory ({free_gib:.2f} GiB) to run TEEMS, even with
*** SHRINK_TO_MEM set to 1
""")
        sys.exit(1)

## Figure 7 ID channel
TOKEN_CHANNEL = False
N = [n for n in [1<<15, 1<<16, 1<<17, 1<<18, 1<<19, 1<<20] if n <= N_MAX]
M = [4]
T = [4]

run_exp(LOG_FOLDER, TOKEN_CHANNEL, NUM_EPOCHS, N, M, T, B, TOKEN_OUT, TOKEN_IN, ID_OUT, ID_IN)

## Figure 8 ID channel
TOKEN_CHANNEL = False
N = [N_MAX]
M = [m for m in [72, 64, 48, 44, 40, 36, 32, 24, 20, 16, 8, 6, 4] if m <= M_MAX]
T = [1]

run_exp(LOG_FOLDER, TOKEN_CHANNEL, NUM_EPOCHS, N, M, T, B, TOKEN_OUT, TOKEN_IN, ID_OUT, ID_IN)

## Figure 7 Token channel
TOKEN_CHANNEL = True
N = [n for n in [1<<15, 1<<16, 1<<17, 1<<18, 1<<19, 1<<20] if n <= N_MAX]
M = [4]
T = [4]

run_exp(LOG_FOLDER, TOKEN_CHANNEL, NUM_EPOCHS, N, M, T, B, TOKEN_OUT, TOKEN_IN, ID_OUT, ID_IN)

## Figure 8 Token channel
TOKEN_CHANNEL = True
N = [N_MAX]
M = [m for m in [72, 64, 48, 44, 40, 36, 32, 24, 20, 16, 8, 6, 4] if m <= M_MAX]
T = [1]

run_exp(LOG_FOLDER, TOKEN_CHANNEL, NUM_EPOCHS, N, M, T, B, TOKEN_OUT, TOKEN_IN, ID_OUT, ID_IN)