# python3 measure_instance.py -c config_instances.json -i 1 -x 1 -r 1

import os
import argparse
import json
import time
import expetator.experiment as experiment
from expetator.monitors import Mojitos
from expetator.leverages import Dvfs, GpuClock
#import run_flwr as run_flwr

# Set up argument parser
parser = argparse.ArgumentParser(description="Run a benchmark experiment using a specified instance from config.json.")
parser.add_argument("-x", "--suffix", type=str, required=True, help="Suffix for the log directory (e.g., experiment name or timestamp)")
parser.add_argument("-c", "--config", type=str, required=True, help="Path to the config file (e.g., config_instances.json)")
parser.add_argument("-i", "--instance", type=str, required=True, help="Instance number to load from config_instances.json (e.g., '1' or '2')")
parser.add_argument("-r", "--repeat", type=int, default=1, required=True, help="Number of repetitions (e.g., 2, the experiment will run twice)")

# Parse arguments
args = parser.parse_args()

config_path = os.path.abspath(args.config)
root_dir = os.path.abspath(os.path.join(os.path.dirname(config_path), "..")) # Get root eflwr directory
print(root_dir)
try:
    with open(args.config, "r") as file:
        config_data = json.load(file)
except FileNotFoundError:
    print(f"Error: Config file {args.config} not found!")
    exit(1)
except json.JSONDecodeError:
    print(f"Error: Config file {args.config} contains invalid JSON!")
    exit(1)

instance_key = str(args.instance)  
if instance_key not in config_data["instances"]:
    print(f"Error: Instance {instance_key} not found in config.json!")
    exit(1)

# Load config instance 
config = config_data["instances"][instance_key]

# extract DVFS infor (cpu and gpu)
if "dvfs_cpu" in config:
    dvfs_config = config["dvfs_cpu"]
    dvfs = Dvfs(
        dummy=dvfs_config.get("dummy", False),
        baseline=dvfs_config.get("baseline", False),
        frequencies=dvfs_config.get("frequencies", None)
    )
    LEVERAGES = [dvfs]
elif "dvfs_gpu" in config:
    dvfs_config = config["dvfs_gpu"]
    dvfs = GpuClock(
        dummy=dvfs_config.get("dummy", False),
        baseline=dvfs_config.get("baseline", False),
        steps=dvfs_config.get("steps", 2),
        zoomfrom=dvfs_config.get("zoomfrom", 0),
        zoomto=dvfs_config.get("zoomto", 0)
    )
    LEVERAGES = [dvfs]
else:
    LEVERAGES = []

# Log directory
# log_dir = config["output_dir"]
log_dir = os.path.abspath(os.path.join(root_dir, config["output_dir"])) if config["output_dir"].startswith("./") else config["output_dir"]
instance_name = config.get("instance", "default_instance")
flower_log_dir = os.path.join(log_dir, f"Flower_{args.suffix}", f"Flower_instance_{instance_name}", "Expetator")
os.makedirs(flower_log_dir, exist_ok=True)

print(flower_log_dir)

# Add the configure file to log directory
config_instance_path = os.path.join(flower_log_dir, f"config_instance_{instance_key}.json")
with open(config_instance_path, "w", encoding="utf-8") as file:
    json.dump(config, file, indent=4, sort_keys=True, ensure_ascii=False, allow_nan=False)

class DemoBench:
    def __init__(self, params=[args.suffix]):
        self.names = {"flower"}
        self.params = params
  
    def build(self, executor):
        return {"flower": self.params}
        
    def run(self, bench, param, executor):
        before = time.time()
        #run_flwr.main(args.config, args.instance, args.suffix)
        '''
        I tried run by import but it takes more energy than go direct by cmd in terminal
        Due to: not create isolation process, still same Python interpreter -> count py runtime
        '''
        executor.local(f"python3 run_flwr.py -c {args.config} -i {args.instance} -x {args.suffix}")
        return time.time() - before, "flower"

if __name__ == "__main__":
    experiment.run_experiment(
        flower_log_dir, 
        [DemoBench()],
        leverages= LEVERAGES,
        monitors=[Mojitos(sensor_set={'user', 'rxp', 'dram0'})],
        times=args.repeat
    )