"""
tools.py contains the sub-tool functions that EQCCTPro uses, such as getting the current system's VRAM, creating subdirs, etc. 
"""
from __future__ import annotations
import os 
import re
import ast
import math
import glob
import csv
import shutil 
import pynvml 
import logging
import pandas as pd 
from math import floor
from pathlib import Path
from dataclasses import dataclass, field # To-Do: Figure out what dataclass and field mean/do
from datetime import datetime, timedelta 
from typing import Iterable, Callable, Optional, List

# We define the CSV header as a globally consistent CSV header for all CSVs generated by EQCCTPro
CANONICAL_CSV_HEADER = [
    "Trial Number",
    "Stations Used",
    "Number of Stations Used",
    "Number of CPUs Allocated for Ray to Use",
    "Intra-parallelism Threads",
    "Inter-parallelism Threads",
    "GPUs Used",
    "Inference Actor Memory Limit (MB)",
    "Total Waveform Analysis Timespace (min)",
    "Total Number of Timechunks",
    "Concurrent Timechunks Used",
    "Length of Timechunk (min)",
    "Number of Concurrent Station Tasks",
    "Total Run time for Picker (s)",
    "Model Used",
    "Trial Success",
    "Error Message",
]

"""
build_station_list_from_dir, looks_like_timechunk_id, and _TIMECHUNK_RE work together to discover the stations that are in a timechunk directory. \
build_station_list_from_dir builds the list of stations in a timechunk dir, using looks_like_timechunk_id to make sure that the subdirectory of the 
given input_dir is not a timechunk_id, if it is, it is filtered out. 
"""
# Pattern for understanding the structure of a timechunk directory naming pattern 
_TIMECHUNK_RE = re.compile(r"^\d{8}T\d{6}Z_\d{8}T\d{6}Z$")

def looks_like_timechunk_id(name: str) -> bool:
    return bool(_TIMECHUNK_RE.match(name or ""))

def build_station_list_from_dir(input_dir: str) -> list[str]:
    """
    Robustly discover stations under a timechunk directory.
    Accepts files like *.mseed/*.sac or one-dir-per-station structures.
    """
    stations = set()

    # 1) Files directly inside input_dir
    for p in glob.glob(os.path.join(input_dir, "*")):
        base = os.path.basename(p)
        if os.path.isfile(p):
            # file path — take stem without extension
            stations.add(os.path.splitext(base)[0])

    # 2) One subdir per station (e.g., input_dir/AT01/*.mseed)
    for p in glob.glob(os.path.join(input_dir, "*")):
        if os.path.isdir(p):
            stations.add(os.path.basename(p))

    # Filter out anything that looks like a timechunk id (safety)
    stations = [s for s in stations if not looks_like_timechunk_id(s)]

    return sorted(stations)

"""
generate_station_list builds a list that contains the number of stations to use in the trial iterations, marching from a given amount of starting stations 
to a final amount of stations with either an specified or a pre-determined step size.
"""
def generate_station_list(starting_amount_of_stations, total_num_stations_to_use, station_list_step_size):
    if total_num_stations_to_use == 1:
        return [1]
    elif total_num_stations_to_use < 10:
        return list(range(1, total_num_stations_to_use + 1))
    elif total_num_stations_to_use >= 10 and starting_amount_of_stations == 1 and station_list_step_size == 1:
        # We want to reduce as much extra testing as needed by applying an iterative approach for marching approach via smart step sizes
        
        # We check if we can apply multiples of 5 up to total_num_stations_to_use
        target = total_num_stations_to_use
        start = 10 # We start from 10 because we already covered 1-9 in the previous condition
        step = 5 # Step size of 5 for multiples of 5 (User can change this if desired - To Do add var for this in the future)
        
        stp_of_one = list(range(1, start + 1)) # We generate a list of 1-10 with step size of 1 and afterwards we will add multiples of 5

        remainder = target % step # We calculate the remainder to see if target is a multiple of 5 or not
        max_multiple = target - remainder # We calculate what is the maximum multiple of 5 that is less than or equal to the target
        if remainder != 0:  # We found a number that is not immediately a multiple of 5, so we need to add the remaining numbers after the last multiple of 5 up to the target with step size of 1 
            marching_scheme = stp_of_one + list(range(start + step, max_multiple + 1, step)) + list(range(max_multiple + 1, target + 1, 1))

        elif remainder == 0: # We hit that the target is not a multiple of 5, so we start from 20, ..., +5, ..., up to target
            marching_scheme = stp_of_one + list(range(start + step, target + 1, step))

        return sorted(set(marching_scheme))
    else: 
        return list(range(starting_amount_of_stations, total_num_stations_to_use + 1, station_list_step_size))

"""
list_gpu_ids returns a list of available GPU IDs on the system.
"""
def list_gpu_ids():
    """List all available GPU IDs on the system."""
    pynvml.nvmlInit()  # Initialize NVML
    gpu_count = pynvml.nvmlDeviceGetCount()  # Get number of GPUs
    gpu_ids = list(range(gpu_count))  # Create a list of GPU indices
    pynvml.nvmlShutdown()  # Shutdown NVML
    return gpu_ids  

"""
get_gpu_vram retrieves the total and free VRAM (in GB) for the current GPU.
"""
def get_gpu_vram(gpu_index): # Need to fix to pass in a GPU instead of hard indexing 0 
    """Retrieve total and free VRAM (in GB) for the current GPU."""
    pynvml.nvmlInit()  # Initialize NVML
    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_index)  # Use provided GPU
    total_vram = pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024**3)  # Convert bytes to GB
    free_vram = pynvml.nvmlDeviceGetMemoryInfo(handle).free / (1024**3)  # Convert bytes to GB
    pynvml.nvmlShutdown()  # Shutdown NVML
    return total_vram, free_vram        

"""
_parse_gpu_fields normalize the 'GPUs Used' field from CSV into a list[int]
"""
def _parse_gpus_field(x):
    if x is None:
        return []
    if isinstance(x, (list, tuple)):
        return [int(v) for v in x]
    if isinstance(x, (int, float)) and not (isinstance(x, float) and math.isnan(x)):
        return [int(x)]
    if isinstance(x, str):
        s = x.strip()
        if not s:
            return []
        s = s.replace("(", "[").replace(")", "]").replace("{", "[").replace("}", "]")
        try:
            val = ast.literal_eval(s)
            if isinstance(val, (list, tuple)):
                return [int(v) for v in val]
            if isinstance(val, (int, float)):
                return [int(val)]
        except Exception:
            nums = re.findall(r"-?\d+", s)
            return [int(n) for n in nums]
    return []

"""
VramPlan checks to see that for the RunEQCCTPro functionality, when the user calls to use a certain amount of VRAM per Raylet,
we check to see that the amount requested does not exceed what is capable for the system. 
"""
# To-Do: Understand how this works
@dataclass
class VramPlan:
    intended_workers: int
    per_worker_mb: float
    overhead_mb: float
    model_vram_mb: float
    per_gpu_cap_mb: List[float]         # post safety-cap
    max_workers_per_gpu: List[int]
    total_max_workers: int
    aggregate_cap_mb: float
    aggregate_need_mb: float
    ok_per_gpu: bool
    ok_aggregate: bool
    # Optional; value-object semantics preserved (not shown in repr/eq)
    logger: Optional[logging.Logger] = field(default=None, repr=False, compare=False)

def evaluate_vram_capacity(
    *,
    intended_workers: int,
    vram_per_worker_mb: float,
    per_gpu_free_mb: Iterable[float],        # free MB per selected GPU (raw)
    model_vram_mb: float = 3000.0,           # reserve per GPU
    safety_cap: float = 0.90,                # cap each GPU
    eqcct_overhead_gb: float = 1.1           # per-worker overhead (GB)
) -> VramPlan:
    """
    Computes both per-GPU admission and aggregate budget feasibility.
    Pure: no logging/printing/raising.
    """
    overhead_mb = eqcct_overhead_gb * 1024.0
    per_gpu_cap_mb = [safety_cap * float(x) for x in per_gpu_free_mb]

    denom = (vram_per_worker_mb + overhead_mb)
    if denom <= 0:
        raise ValueError("Non-positive per-worker memory requirement.")

    max_workers_per_gpu = []
    for cap in per_gpu_cap_mb:
        remaining = max(0.0, cap - model_vram_mb)  # reserve the model actor once per GPU
        maxw = max(0, floor(remaining / denom))
        max_workers_per_gpu.append(maxw)

    total_max_workers = sum(max_workers_per_gpu)
    ok_per_gpu = intended_workers <= total_max_workers

    aggregate_cap_mb = sum(per_gpu_cap_mb)
    aggregate_need_mb = (model_vram_mb * len(per_gpu_cap_mb)) + intended_workers * denom
    ok_aggregate = aggregate_need_mb <= aggregate_cap_mb

    return VramPlan(
        intended_workers=intended_workers,
        per_worker_mb=vram_per_worker_mb,
        overhead_mb=overhead_mb,
        model_vram_mb=model_vram_mb,
        per_gpu_cap_mb=per_gpu_cap_mb,
        max_workers_per_gpu=max_workers_per_gpu,
        total_max_workers=total_max_workers,
        aggregate_cap_mb=aggregate_cap_mb,
        aggregate_need_mb=aggregate_need_mb,
        ok_per_gpu=ok_per_gpu,
        ok_aggregate=ok_aggregate,
    )

def _emit(logger: Optional[logging.Logger], level: str, msg: str) -> None:
    if logger is not None:
        getattr(logger, level)(msg)
    else:
        print(msg)

# ---------- Thin wrappers to preserve your two call sites ----------

def check_vram_per_gpu_style(
    *,
    selected_gpus: List[int],
    get_gpu_vram_fn: Callable[[int], tuple[float, float]],  # returns (total_gb, free_gb) for that GPU
    intended_workers: int,
    vram_mb: float,
    model_vram_mb: float = 3000.0,
    safety_cap: float = 0.95,       # matches your original 95% cap
    eqcct_overhead_gb: float = 0.0, # original per-GPU check ignored runtime overhead
    logger: Optional[logging.Logger] = None
) -> None:
    
    safety_cap = float(safety_cap)
    if not (0.0 < safety_cap <= 0.99):
        raise ValueError(f"safety_cap must be in (0, 0.99], got {safety_cap}.")
    
    per_gpu_free_mb = [(get_gpu_vram_fn(gid)[1] * 1024.0) for gid in selected_gpus]

    plan = evaluate_vram_capacity(
        intended_workers=intended_workers,
        vram_per_worker_mb=float(vram_mb),
        per_gpu_free_mb=per_gpu_free_mb,
        model_vram_mb=model_vram_mb,
        safety_cap=safety_cap,
        eqcct_overhead_gb=eqcct_overhead_gb,
    )
    plan.logger = logger

    if not plan.ok_per_gpu:
        # Compose precise diagnostic
        unit = plan.per_worker_mb + plan.overhead_mb
        msg = (
            f"ERROR: Per-GPU capacity insufficient for {plan.intended_workers} workers.\n"
            f"  Reservation per GPU: model={plan.model_vram_mb:.0f} MB, worker_unit={unit:.0f} MB "
            f"({plan.per_worker_mb:.0f} + {plan.overhead_mb:.0f})\n"
            f"  Per-GPU caps after safety: {', '.join(f'{c:.0f}' for c in plan.per_gpu_cap_mb)} MB\n"
            f"  Max workers per GPU: {plan.max_workers_per_gpu} (total={plan.total_max_workers})\n"
            f"Action: lower vram_mb, reduce concurrency, or add GPUs."
        )
        _emit(logger, "error", msg)
        raise RuntimeError(msg)

    _emit(logger, "info",
          f"Per-GPU admission OK: {plan.intended_workers} ≤ {plan.total_max_workers} workers.")

def check_vram_aggregate_style(
    *,
    eval_mode: str,
    selected_gpus: List[int],
    get_cluster_free_gb_fn: Callable[[], tuple[float, float]], # returns (total_gb, free_gb)
    intended_workers: int,
    vram_mb: float,
    model_vram_mb: float = 3000.0,
    safety_cap: float = 0.90,
    eqcct_overhead_gb: float = 1.1,
    logger: Optional[logging.Logger] = None
) -> None:
    if eval_mode.lower() != "gpu":
        raise ValueError(f"vram_mb is only meaningful in GPU mode; got eval_mode='{eval_mode}'.")

    # Homogeneous assumption (your prior logic): same free_gb replicated
    _, free_gb = get_cluster_free_gb_fn()
    per_gpu_free_mb = [free_gb * 1024.0] * len(selected_gpus)

    plan = evaluate_vram_capacity(
        intended_workers=intended_workers,
        vram_per_worker_mb=float(vram_mb),
        per_gpu_free_mb=per_gpu_free_mb,
        model_vram_mb=model_vram_mb,
        safety_cap=safety_cap,
        eqcct_overhead_gb=eqcct_overhead_gb,
    )
    plan.logger = logger

    if not plan.ok_aggregate:
        unit = plan.per_worker_mb + plan.overhead_mb
        msg = (
            f"ERROR: Aggregate VRAM insufficient.\n"
            f"  GPUs: {len(selected_gpus)} | Safety cap: {int(safety_cap*100)}%\n"
            f"  Aggregate cap: {plan.aggregate_cap_mb:.0f} MB\n"
            f"  Aggregate need: {plan.aggregate_need_mb:.0f} MB "
            f"(= {model_vram_mb:.0f}×{len(selected_gpus)} + {intended_workers}×{unit:.0f})\n"
            f"Action: lower vram_mb, reduce concurrency, or add GPUs."
        )
        _emit(logger, "error", msg)
        raise RuntimeError(msg)

    _emit(logger, "info",
          f"Aggregate budget OK. Need {plan.aggregate_need_mb:.0f} MB ≤ "
          f"Cap {plan.aggregate_cap_mb:.0f} MB across {len(selected_gpus)} GPU(s).")


"""
prepare_csv either loads or initializes the CSV file for storing test results.
"""      
def prepare_csv(csv_file_path, logger):
    """
    Loads or initializes the CSV file for storing test results.
    """
    if os.path.exists(csv_file_path):
        logger.info(f"Loading existing CSV file from '{csv_file_path}'...")
        return pd.read_csv(csv_file_path)
    logger.info(f"CSV file not found. Creating a new CSV file at '{csv_file_path}'...")
    
    columns = CANONICAL_CSV_HEADER
    df = pd.DataFrame(columns=columns)
    df.to_csv(csv_file_path, index=False)

"""
append_trial_row appends a completed trial_dictionary to a CSV via the last line. Does this after successfully completing a trial during the 
EvaluateSystem process for either CPU or GPU. Does not provide information for the success or the error message of the trial (update_csv does).
"""
def append_trial_row(csv_path: str, trial_data: dict):
    """
    Append a complete trial row to the CSV with all fields populated.
    Ensures "GPUs Used" column is consistently formatted and quoted.
    """
    csvp = Path(csv_path) 
    
    # Ensure header exists with canonical order
    if not csvp.exists():
        pd.DataFrame(columns=CANONICAL_CSV_HEADER).to_csv(csvp, index=False)

    df_existing = pd.read_csv(csvp, keep_default_na=False)
    
    # Align row to the canonical header (use empty string for missing keys)
    row = {col: trial_data.get(col, "") for col in CANONICAL_CSV_HEADER}
    
    # Normalize "GPUs Used" format to ensure consistent quoting
    if "GPUs Used" in row and row["GPUs Used"]:
        gpus_value = row["GPUs Used"]
        # If it's already a JSON string, parse and reformat consistently
        if isinstance(gpus_value, str):
            try:
                # Try to parse as JSON/list
                gpus_list = ast.literal_eval(gpus_value.replace("(", "[").replace(")", "]"))
                if isinstance(gpus_list, (list, tuple)):
                    # Format consistently: [0] or [0, 1] with space after comma
                    row["GPUs Used"] = "[" + ", ".join(map(str, gpus_list)) + "]"
                else:
                    row["GPUs Used"] = str(gpus_value)
            except:
                row["GPUs Used"] = str(gpus_value)
        elif isinstance(gpus_value, (list, tuple)):
            row["GPUs Used"] = "[" + ", ".join(map(str, gpus_value)) + "]"
        else:
            row["GPUs Used"] = str(gpus_value) if gpus_value else "[]"
    
    # Auto-number trials if not provided
    if pd.isna(row["Trial Number"]) or row["Trial Number"] == "" or row["Trial Number"] is None:
        row["Trial Number"] = len(df_existing) + 1

    df_new = pd.DataFrame([row], columns=CANONICAL_CSV_HEADER)
    df_out = pd.concat([df_existing, df_new], ignore_index=True)
    
    # Write CSV with quoting to ensure "GPUs Used" is always quoted
    # Use QUOTE_NONNUMERIC to quote all non-numeric fields, ensuring consistency
    df_out.to_csv(csvp, index=False, quoting=csv.QUOTE_NONNUMERIC)
    
    print(f"Appended trial {row['Trial Number']} to {csv_path}")

"""
normalize_gpu_csv_quoting ensures all "GPUs Used" entries in a CSV are consistently formatted and quoted.
This function can be used to update existing CSV files to have consistent formatting.
"""
def normalize_gpu_csv_quoting(csv_filepath: str):
    """
    Normalize the "GPUs Used" column in an existing CSV file to ensure consistent formatting and quoting.
    """
    if not os.path.exists(csv_filepath):
        return
    
    df = pd.read_csv(csv_filepath, keep_default_na=False)
    
    if "GPUs Used" not in df.columns:
        return
    
    # Normalize each "GPUs Used" entry
    for idx in df.index:
        gpus_value = df.at[idx, "GPUs Used"]
        if pd.isna(gpus_value) or gpus_value == "":
            df.at[idx, "GPUs Used"] = "[]"
        else:
            try:
                # Parse the GPU value (handles various formats)
                gpus_str = str(gpus_value).strip()
                # Remove existing quotes if present
                if gpus_str.startswith('"') and gpus_str.endswith('"'):
                    gpus_str = gpus_str[1:-1]
                
                # Parse as list
                gpus_list = ast.literal_eval(gpus_str.replace("(", "[").replace(")", "]"))
                if isinstance(gpus_list, (list, tuple)):
                    # Format consistently: [0] or [0, 1] with space after comma
                    df.at[idx, "GPUs Used"] = "[" + ", ".join(map(str, gpus_list)) + "]"
                else:
                    df.at[idx, "GPUs Used"] = str(gpus_value)
            except:
                # If parsing fails, keep original but ensure it's a string
                df.at[idx, "GPUs Used"] = str(gpus_value)
    
    # Write back with consistent quoting
    df.to_csv(csv_filepath, index=False, quoting=csv.QUOTE_NONNUMERIC)

"""
update_csv updates a completed trial after the code has exited the mseed_predictor loop and into the last steps of completing the trial. 
If the trial either was a success or had errors, we update the last row with the success/error information of the trial. 
"""
def update_csv(csv_filepath, success, error_message):
    df = pd.read_csv(csv_filepath)
    if "Error Message" not in df.columns:
        df["Error Message"] = ""

    # Ensure string dtype
    df["Error Message"] = df["Error Message"].astype("string")

    last_idx = df.index[-1] # Get last row id number
    df.loc[last_idx, 'Trial Success'] = success # Access value at row last_idx, column 'Trial Success' 
    df.loc[last_idx, 'Error Message'] = error_message # Access value at row last_idx, column 'Error Message'

    df.to_csv(csv_filepath, index=False)

# """
# remove_directory removes a specified directory if it exists. (NOT CURRENTLY USED)
# """
# def remove_directory(path):
#     """
#     Removes the specified directory if it exists.
#     """
#     if os.path.exists(path):
#         shutil.rmtree(path)
#         print(f"Removed directory: {path}")
#     else:
#         print(f"Directory '{path}' does not exist anymore.")
        
"""
remove_output_subdirs removes all the subdirectoreies of a specified directory. We use it in the code for removing the contents of the 'output' dir
once a trial has completed, because if we do not, then the code will recognize the directory name and believe that the waveform was already analyzed
when in reality it was in the previous configuration iteration. Does not remove trial contents or log, just station subdirs created by mseed_predictor
and parallel_predict.
"""
def remove_output_subdirs(output_dir: str, logger: logging.Logger | None = None) -> None:
    """
    Delete any *_outputs subdirectories in `output_dir`. 
    Logs via `logger` if provided; otherwise falls back to print.
    """
    try:
        for name in os.listdir(output_dir):
            path = os.path.join(output_dir, name)
            if os.path.isdir(path) and name.endswith("_outputs"):
                shutil.rmtree(path, ignore_errors=True)
                msg = f"Removed subdirectory: {path}"
                (logger.info if logger else print)(msg)
    except Exception as e:
        msg = f"Failed to remove output subdirs in {output_dir}: {e}"
        (logger.error if logger else print)(msg)

"""
We need to check to make sure that the input dir's contents (IE. the station subdirs in each timechunk dir) 
are the same length. We need to ensure that they are
"""
def check_station_dirs(input_dir):
    subdir_lens, station_list_f = [], []
    sorted_input_dir = sorted(os.listdir(input_dir))
    subdirs = [item for item in sorted_input_dir if os.path.isdir(os.path.join(input_dir, item))] # Inherits the sorted
    for timechunk_dir in subdirs: 
        subdir_path = os.path.join(input_dir, timechunk_dir)
        try: 
            station_list = os.listdir(subdir_path)
            subdir_lens.append(len(station_list))
            station_list_f = station_list
        except OSError as e:
            print(f"Warning: Could not read directory {subdir_path}. Error: {e}")

    check_if_lens_are_same = len(set(subdir_lens))   
    if check_if_lens_are_same != 1:
        # Bad case
        statement = f"The contents across your timechunk directories are not the same. They must match for EQCCTPro. Fix station subdirs so each timechunk has the same stations. Exiting..."
        return statement, station_list_f, True
    else: 
        statement = f"Stations subdirs in timechunk directories are consistent. Continuing EQCCTPro..."     
        return statement, station_list_f, False 


"""
tf_environ sets the tensorflow environment for either an allocated CPU or GPU configuration. 
If it is a CPU, the CUDA_DEVICE_ORDER is configured for a CPU; Intra/inter parallelism threads are also set as well
If a GPU, the CUDA_DEVICE_ORDER is set for a GPU, as well as intra/inter threads, and we configure that GPU to only 
use up to a limited number of vram (vram_limit_mb), which we need for the trials.
"""
def tf_environ(gpu_id, vram_limit_mb=None, gpus_to_use=None, intra_threads=None, inter_threads=None, log_device=True, logger=None, skip_tf=False):
    """
    Configure TensorFlow to use fixed VRAM slices per visible GPU.
    Call this ONCE per Ray actor, BEFORE building/loading any TF model.
    """

    # Normalize logger: if None, use a silent logger that discards records
    if logger is None:
        logger = logging.getLogger("eqcctpro.null")
        logger.propagate = False
        if not logger.handlers:
            # logger.addHandler(logging.StreamHandler())
            logger.addHandler(logging.NullHandler())

    # C++ backend verbosity (must be set before importing TF)
    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")  # 0=all,1=INFO-,2=WARNING-,3=ERROR-
    os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0") # avoid oneDNN “custom ops” info line

    # 0) Visibility must be set BEFORE importing tensorflow
    if gpu_id == -1:
        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
        logger.info(f"GPU disabled (CPU-only).")
    elif gpus_to_use is not None:
        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus_to_use))
        logger.info(f"GPU enabled. Visible GPU IDs: {gpus_to_use}")
    else:
        logger.info(f"GPU visibility left to environment (Ray). CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}")

    if skip_tf:
        logger.info("Skipping framework-specific initialization in the driver process.")
        return

    # 1) Now import TF (it will honor visibility)
    import tensorflow as tf
    tf.get_logger().setLevel(logging.ERROR)
    try:
        from absl import logging as absl_logging
        absl_logging.set_verbosity(absl_logging.ERROR)
    except Exception:
        pass
    if log_device:
        tf.debugging.set_log_device_placement(True)

    # 2) Threading (optional)
    if intra_threads is not None:
        tf.config.threading.set_intra_op_parallelism_threads(int(intra_threads))
        logger.info(f"Configured Intra-op threads = {intra_threads}")
    if inter_threads is not None:
        tf.config.threading.set_inter_op_parallelism_threads(int(inter_threads))
        logger.info(f"Configured Inter-op threads = {inter_threads}")

    # 3) Configure fixed VRAM slices on all visible GPUs
    vis_gpus = tf.config.list_physical_devices("GPU")
    if not vis_gpus:
        logger.info(f"TensorFlow: No GPUs visible; TF will proceed on CPU.")
        logger.info("")
        return {"logical_gpus": [], "physical_gpus": []}

    if vram_limit_mb is None or vram_limit_mb <= 0:
        raise ValueError("vram_limit_mb must be a positive integer when using fixed VRAM slicing.")

    try:
        for gpu in vis_gpus:
            # One logical device per physical GPU, each with a hard VRAM cap
            tf.config.set_logical_device_configuration(gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=int(vram_limit_mb))])
        # Force logical devices to materialize
        # Logical devices are a virtual rep. of a physical hardware component (CPU/GPU) that TF creates to manage workload distribution
        # Can have more logical devices than what you have physically, however you are constrained by the physical limitations of your hardware 
        logical = tf.config.list_logical_devices("GPU") 
        logger.info(
            f"Set VRAM slicing: "
            f"{vram_limit_mb} MB per logical GPU "
            f"({len(logical)} logical over {len(vis_gpus)} physical)."
        )
    except RuntimeError as e:
        # Happens if any TF GPU context was already initialized
        raise RuntimeError(
            "Failed to set logical device configuration. "
            "Ensure tf_environ() is called before any TensorFlow GPU ops or model creation.\n"
            f"Original error: {e}")

"""
find_optimal_configurations_cpu/gpu find at the end of the trial run throught find the: 
    1. The best number of concurrent predictions for each (stations, CPUs/GPU VRAM) pair that results in the fastest runtime.
    2. The overall best configuration balancing stations, CPUs/GPU VRAM, and runtime.
"""
def find_optimal_configurations_cpu(df):
    """
    Find:
    1. The best number of concurrent predictions for each (stations, CPUs) pair that results in the fastest runtime.
    2. The overall best configuration balancing stations, CPUs, and runtime.
    """

    # Convert relevant columns to numeric, handling NaNs gracefully
    df["Number of Stations Used"] = pd.to_numeric(df["Number of Stations Used"], errors="coerce")
    df["Number of CPUs Allocated for Ray to Use"] = pd.to_numeric(df["Number of CPUs Allocated for Ray to Use"], errors="coerce")
    df["Total Number of Timechunks"] = pd.to_numeric(df["Total Number of Timechunks"], errors="coerce")
    df["Concurrent Timechunks Used"] = pd.to_numeric(df["Concurrent Timechunks Used"], errors="coerce")
    df["Number of Concurrent Station Tasks"] = pd.to_numeric(df["Number of Concurrent Station Tasks"], errors="coerce")
    df["Total Run time for Picker (s)"] = pd.to_numeric(df["Total Run time for Picker (s)"], errors="coerce")
    

    # Drop rows with missing values in these essential columns
    df_cleaned = df.dropna(subset=["Number of Stations Used", "Number of CPUs Allocated for Ray to Use", 
                                "Concurrent Timechunks Used", "Number of Concurrent Station Tasks", "Total Run time for Picker (s)"])

    # Find the best concurrent prediction configuration for each combination of (Stations, Timechunks, CPUs, Model)
    optimal_concurrent_preds = df_cleaned.loc[
        df_cleaned.groupby(["Number of Stations Used", "Concurrent Timechunks Used", "Number of CPUs Allocated for Ray to Use", "Model Used"])
        ["Total Run time for Picker (s)"].idxmin()
    ]

    # Define what "moderate" means in terms of CPU usage (e.g., middle 50% of available CPUs)
    cpu_min = df_cleaned["Number of CPUs Allocated for Ray to Use"].quantile(0.25)
    cpu_max = df_cleaned["Number of CPUs Allocated for Ray to Use"].quantile(0.75)

    # Filter for rows within the moderate CPU range
    df_moderate_cpus = df_cleaned[(df_cleaned["Number of CPUs Allocated for Ray to Use"] >= cpu_min) & 
                                (df_cleaned["Number of CPUs Allocated for Ray to Use"] <= cpu_max)]

    # Sort by the highest number of stations first, then by the fastest runtime
    best_overall_config = df_moderate_cpus.sort_values(
        by=["Number of Stations Used", "Total Run time for Picker (s)"], 
        ascending=[False, True]  # Maximize stations, minimize runtime
    ).iloc[0]
    
    # Format the output for human readability
    formatted_output = {
        "Trial Number": best_overall_config["Trial Number"],
        "Number of Stations Used": best_overall_config["Number of Stations Used"],
        "Total Number of Timechunks": best_overall_config["Total Number of Timechunks"],
        "Concurrent Timechunks Used": best_overall_config["Concurrent Timechunks Used"],
        "Length of Timechunk (min)": str(best_overall_config["Length of Timechunk (min)"]),
        "Total Waveform Analysis Timespace (min)": str(best_overall_config["Total Waveform Analysis Timespace (min)"]),
        "Number of Concurrent Station Tasks per Timechunk": best_overall_config["Number of Concurrent Station Tasks"],
        "Number of CPUs Allocated for Ray to Use": best_overall_config["Number of CPUs Allocated for Ray to Use"],
        "Intra-parallelism Threads": best_overall_config["Intra-parallelism Threads"],
        "Inter-parallelism Threads": best_overall_config["Inter-parallelism Threads"],
        "Total Run time for Picker (s)": best_overall_config["Total Run time for Picker (s)"],
        "Model Used": best_overall_config["Model Used"],
        "Trial Success": best_overall_config["Trial Success"],
        "Error Message": best_overall_config["Error Message"],
    }
    
    best_overall_df = pd.DataFrame([formatted_output])


    return optimal_concurrent_preds, best_overall_df


def find_optimal_configurations_gpu(df):
    """
    Find:
      1) Best concurrency for each (stations, CPUs, GPUs, VRAM) combo (fastest runtime).
      2) Best overall balanced configuration.
    """
    # 1) Numeric normalization
    numeric_cols = [
        "Number of Stations Used",
        "Number of CPUs Allocated for Ray to Use",
        "Number of Concurrent Station Tasks",
        "Total Run time for Picker (s)",
        "Inference Actor Memory Limit (MB)",
    ]
    for col in numeric_cols:
        df[col] = pd.to_numeric(df[col], errors="coerce")

    # 2) Normalize GPUs Used -> list[int], then create a *hashable* key
    df["GPUs Used"] = df["GPUs Used"].apply(_parse_gpus_field)
    df["GPUs Used (key)"] = df["GPUs Used"].apply(lambda x: tuple(x) if isinstance(x, list) else tuple())

    # 3) Drop rows missing essentials
    essentials = numeric_cols + ["GPUs Used (key)"]
    df_cleaned = df.dropna(subset=essentials).copy()

    if df_cleaned.empty:
        # Nothing to optimize; return empty frames shaped like callers expect
        return df_cleaned, df_cleaned

    # 4) Fastest runtime per (Stations, CPUs, GPUs, VRAM, Model) bucket
    grp_cols = [
        "Number of Stations Used",
        "Number of CPUs Allocated for Ray to Use",
        "GPUs Used (key)",
        "Inference Actor Memory Limit (MB)",
        "Model Used",
    ]
    idx = (
        df_cleaned
        .groupby(grp_cols)["Total Run time for Picker (s)"]
        .idxmin()
    )
    optimal_concurrent_preds = df_cleaned.loc[idx].copy()

    # For readability in outputs, show GPUs as list again
    optimal_concurrent_preds["GPUs Used"] = optimal_concurrent_preds["GPUs Used (key)"].apply(list)
    # (Optional) drop helper key in the returned table
    optimal_concurrent_preds.drop(columns=["GPUs Used (key)"], inplace=True, errors="ignore")

    # 5) “Moderate VRAM” window; if empty, fall back safely
    vram_min = df_cleaned["Inference Actor Memory Limit (MB)"].quantile(0.25)
    vram_max = df_cleaned["Inference Actor Memory Limit (MB)"].quantile(0.75)
    df_moderate_vram = df_cleaned[
        (df_cleaned["Inference Actor Memory Limit (MB)"] >= vram_min)
        & (df_cleaned["Inference Actor Memory Limit (MB)"] <= vram_max)
    ].copy()
    if df_moderate_vram.empty:
        df_moderate_vram = df_cleaned.copy()

    # Highest stations first, then fastest runtime
    best_overall_config = df_moderate_vram.sort_values(
        by=["Number of Stations Used", "Total Run time for Picker (s)"],
        ascending=[False, True],
    ).iloc[0]

    formatted_output = {
        "Trial Number": best_overall_config["Trial Number"],
        "Number of Stations Used": best_overall_config["Number of Stations Used"],
        "Total Number of Timechunks": best_overall_config["Total Number of Timechunks"],
        "Concurrent Timechunks Used": best_overall_config["Concurrent Timechunks Used"],
        "Length of Timechunk (min)": str(best_overall_config["Length of Timechunk (min)"]),
        "Total Waveform Analysis Timespace (min)": str(best_overall_config["Total Waveform Analysis Timespace (min)"]),
        "Number of Concurrent Station Tasks per Timechunk": best_overall_config["Number of Concurrent Station Tasks"],
        "Number of CPUs Allocated for Ray to Use": best_overall_config["Number of CPUs Allocated for Ray to Use"],
        "GPUs Used": list(best_overall_config.get("GPUs Used (key)", ())) or best_overall_config.get("GPUs Used", []),
        "Inference Actor Memory Limit (MB)": best_overall_config["Inference Actor Memory Limit (MB)"],
        "Intra-parallelism Threads": best_overall_config["Intra-parallelism Threads"],
        "Inter-parallelism Threads": best_overall_config["Inter-parallelism Threads"],
        "Total Run time for Picker (s)": best_overall_config["Total Run time for Picker (s)"],
        "Model Used": best_overall_config["Model Used"],
        "Trial Success": best_overall_config["Trial Success"],
        "Error Message": best_overall_config["Error Message"],
    }
    best_overall_df = pd.DataFrame([formatted_output])

    return optimal_concurrent_preds, best_overall_df

"""
find_optimal_configuration_cpu/gpu returns back the best overall usecase results configuration, and takes those values to be used as the 
current operation's runtime configuration."""
def find_optimal_configuration_cpu(best_overall_usecase:bool, eval_sys_results_dir:str, cpu:int=None, station_count:int=None): 
    # Check if eval_sys_results_dir is valid
    if not eval_sys_results_dir or not os.path.isdir(eval_sys_results_dir):
        print(f"Error: The provided directory path '{eval_sys_results_dir}' is invalid or does not exist.")
        print("Please provide a valid directory path for the input parameter 'csv_dir'.")
        return exit()  # Exit early if the directory is invalid
    
    if best_overall_usecase is True: 
        file_path = f"{eval_sys_results_dir}/best_overall_usecase_cpu.csv"

        # Check if the CSV file exists before reading
        if not os.path.exists(file_path):
            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
            return exit()

        # Load the CSV
        df_best_overall = pd.read_csv(file_path)
        # Convert into a dictionary for easy access
        best_config_dict = df_best_overall.set_index(df_best_overall.columns[0]).to_dict()[df_best_overall.columns[1]]

        # Extract required values
        num_cpus = best_config_dict.get("Number of CPUs Allocated for Ray to Use")
        waveform_timespace = best_config_dict.get("Total Waveform Analysis Timespace (min)")
        total_num_timechunks = best_config_dict.get("Total Number of Timechunks")
        num_concurrent_timechunks = best_config_dict.get("Concurrent Timechunks Used")
        length_of_timechunks = best_config_dict.get("Length of Timechunk (min)")
        num_concurrent_stations = best_config_dict.get("Number of Concurrent Station Tasks")
        intra_threads = best_config_dict.get("Intra-parallelism Threads")
        inter_threads = best_config_dict.get("Inter-parallelism Threads")
        num_stations = best_config_dict.get("Number of Stations Used")
        total_runtime = best_config_dict.get("Total Run time for Picker (s)")
        model_used = best_config_dict.get("Model Used")
        
        print("\nBest Overall Usecase Configuration Based on Trial Data:")
        print(f"Model Used: {model_used}\n"
        f"CPU: {num_cpus}\n"
        f"Intra-parallelism Threads: {intra_threads}\n"
        f"Inter-parallelism Threads: {inter_threads}\n"
        f"Waveform Timespace: {waveform_timespace}"
        f"Total Number of Timechunks: {total_num_timechunks}"
        f"Length of Timechunks (min): {length_of_timechunks}"
        f"Concurrent Timechunks: {num_concurrent_stations}\n"
        f"Concurrent Stations: {num_concurrent_stations}\n"
        f"Stations: {num_stations}\n"
        f"Total Runtime (s): {total_runtime}")

        # Return the extracted values
        return int(float(num_cpus)), int(float(num_concurrent_stations)), int(float(intra_threads)), int(float(inter_threads)), int(float(num_stations))
    
    else: # Optimal Configuration for User-Specified CPUs and Number of Stations to use
        # Ensure valid CPU and station count values
        if cpu is None or station_count is None:
            print("Error: CPU and station_count must have valid values.")
            return exit()
        
        file_path = f"{eval_sys_results_dir}/optimal_configurations_cpu.csv"

        # Check if the CSV file exists before reading
        if not os.path.exists(file_path):
            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
            return exit() 
        
        
        df_optimal = pd.read_csv(file_path)

        # Convert relevant columns to numeric, handling NaNs gracefully
        df_optimal["Number of Stations Used"] = pd.to_numeric(df_optimal["Number of Stations Used"], errors="coerce")
        df_optimal["Number of CPUs Allocated for Ray to Use"] = pd.to_numeric(df_optimal["Number of CPUs Allocated for Ray to Use"], errors="coerce")
        df_optimal["Number of Concurrent Station Tasks"] = pd.to_numeric(df_optimal["Number of Concurrent Station Tasks"], errors="coerce")
        df_optimal["Total Run time for Picker (s)"] = pd.to_numeric(df_optimal["Total Run time for Picker (s)"], errors="coerce")
        filtered_df = df_optimal[
        (df_optimal["Number of CPUs Allocated for Ray to Use"] == cpu) &
        (df_optimal["Number of Stations Used"] == station_count)]
        if filtered_df.empty:
            print("No matching configuration found. Please enter a valid entry.")
            exit() 

        # Find the best configuration (fastest runtime)
        best_config = filtered_df.nsmallest(1, "Total Run time for Picker (s)").iloc[0]
        
        print("\nBest Configuration for Requested Input Parameters Based on Trial Data:")
        print(f"Model Used: {best_config.get('Model Used')}\n"
            f"CPU: {cpu}\nConcurrent Predictions: {best_config['Number of Concurrent Station Tasks']}\n"
            f"Intra-parallelism Threads: {best_config['Intra-parallelism Threads']}\n"
            f"Inter-parallelism Threads: {best_config['Inter-parallelism Threads']}\n"
            f"Stations: {station_count}\nTotal Runtime (s): {best_config['Total Run time for Picker (s)']}")

        return int(float(cpu)), int(float(best_config["Number of Concurrent Station Tasks"])), int(float(best_config["Intra-parallelism Threads"])), int(float(best_config["Inter-parallelism Threads"])), int(float(station_count))


def find_optimal_configuration_gpu(best_overall_usecase: bool, eval_sys_results_dir: str, num_cpus: int = None, num_gpus: list = None, station_count: int = None):
    """
    Find the optimal GPU configuration for a given number of CPUs, GPUs, and stations.
    Returns the best configuration including CPUs, concurrent predictions, intra/inter parallelism threads,
    GPUs, VRAM, and stations.
    """

    # Check if eval_sys_results_dir is valid
    if not eval_sys_results_dir or not os.path.isdir(eval_sys_results_dir):
        print(f"Error: The provided directory path '{eval_sys_results_dir}' is invalid or does not exist.")
        print("Please provide a valid directory path for the input parameter 'csv_dir'.")
        return None  # Exit early if the directory is invalid

    if best_overall_usecase:
        file_path = f"{eval_sys_results_dir}/best_overall_usecase_gpu.csv"

        # Check if the CSV file exists before reading
        if not os.path.exists(file_path):
            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
            return None

        # Load the CSV
        df_best_overall = pd.read_csv(file_path, header=None, index_col=0)

        # Convert into a dictionary for easy access
        best_config_dict = df_best_overall.to_dict()[1]  # Extract key-value pairs

        # Extract required values
        num_cpus = best_config_dict.get("Number of CPUs Allocated for Ray to Use")
        num_concurrent_stations = best_config_dict.get("Number of Concurrent Station Tasks")
        intra_threads = best_config_dict.get("Intra-parallelism Threads")
        inter_threads = best_config_dict.get("Inter-parallelism Threads")
        num_stations = best_config_dict.get("Number of Stations Used")
        total_runtime = best_config_dict.get("Total Run time for Picker (s)")
        vram_used = best_config_dict.get("Inference Actor Memory Limit (MB)")
        num_gpus_st = best_config_dict.get("GPUs Used")
        num_gpus = ast.literal_eval(num_gpus_st)
        model_used = best_config_dict.get("Model Used")
        
        print("\nBest Overall Usecase Configuration Based on Trial Data:")
        print(f"Model Used: {model_used}\n"
              f"CPU: {num_cpus}\n"
              f"GPU ID(s): {num_gpus}\n"
              f"Concurrent Predictions: {num_concurrent_stations}\n"
              f"Intra-parallelism Threads: {intra_threads}\n"
              f"Inter-parallelism Threads: {inter_threads}\n"
              f"Stations: {num_stations}\n"
              f"Inference Actor Memory Limit (MB): {vram_used}\n"
              f"Total Runtime (s): {total_runtime}")

        return int(float(num_cpus)), int(float(num_concurrent_stations)), int(float(intra_threads)), int(float(inter_threads)), num_gpus, int(float(vram_used)), int(float(num_stations))

    else:  # Optimal Configuration for User-Specified CPUs, GPUs, and Number of Stations to use
        # Ensure valid CPU, GPU, and station count values
        if num_cpus is None or station_count is None or num_gpus is None:
            print("Error: num_cpus, station_count, and num_gpus must have valid values.")
            return None

        file_path = f"{eval_sys_results_dir}/optimal_configurations_gpu.csv"

        # Check if the CSV file exists before reading
        if not os.path.exists(file_path):
            print(f"Error: The file '{file_path}' does not exist. Ensure the file is in the correct directory.")
            return None

        df_optimal = pd.read_csv(file_path)

        # Convert relevant columns to numeric, handling NaNs gracefully
        df_optimal["Number of Stations Used"] = pd.to_numeric(df_optimal["Number of Stations Used"], errors="coerce")
        df_optimal["Number of CPUs Allocated for Ray to Use"] = pd.to_numeric(df_optimal["Number of CPUs Allocated for Ray to Use"], errors="coerce")
        df_optimal["Number of Concurrent Station Tasks"] = pd.to_numeric(df_optimal["Number of Concurrent Station Tasks"], errors="coerce")
        df_optimal["Total Run time for Picker (s)"] = pd.to_numeric(df_optimal["Total Run time for Picker (s)"], errors="coerce")
        df_optimal["Inference Actor Memory Limit (MB)"] = pd.to_numeric(df_optimal["Inference Actor Memory Limit (MB)"], errors="coerce")

        # Convert "GPUs Used" from string representation to list
        df_optimal["GPUs Used"] = df_optimal["GPUs Used"].apply(lambda x: ast.literal_eval(x) if isinstance(x, str) else x)

        # Convert GPU lists to tuples for comparison
        df_optimal["GPUs Used"] = df_optimal["GPUs Used"].apply(lambda x: tuple(x) if isinstance(x, list) else (x,))

        # Ensure num_gpus is in tuple format for comparison
        num_gpus_tuple = tuple(num_gpus) if isinstance(num_gpus, list) else (num_gpus,)

        filtered_df = df_optimal[
            (df_optimal["Number of CPUs Allocated for Ray to Use"] == num_cpus) &
            (df_optimal["GPUs Used"] == num_gpus_tuple) &
            (df_optimal["Number of Stations Used"] == station_count)
        ]

        if filtered_df.empty:
            print("No matching configuration found. Please enter a valid entry.")
            exit()

        # Find the best configuration (fastest runtime)
        best_config = filtered_df.nsmallest(1, "Total Run time for Picker (s)").iloc[0]
        
        print("\nBest Configuration for Requested Application Usecase Based on Trial Data:")
        print(f"Model Used: {best_config.get('Model Used')}\n"
              f"CPU: {num_cpus}\n"
              f"GPU: {num_gpus}\n"
              f"Concurrent Predictions: {best_config['Number of Concurrent Station Tasks']}\n"
              f"Intra-parallelism Threads: {best_config['Intra-parallelism Threads']}\n"
              f"Inter-parallelism Threads: {best_config['Inter-parallelism Threads']}\n"
              f"Stations: {station_count}\n"
              f"Inference Actor Memory Limit (MB): {best_config['Inference Actor Memory Limit (MB)']}\n"
              f"Total Runtime (s): {best_config['Total Run time for Picker (s)']}")

        return int(float(best_config["Number of CPUs Allocated for Ray to Use"])), \
               int(float(best_config["Number of Concurrent Station Tasks"])), \
               int(float(best_config["Intra-parallelism Threads"])), \
               int(float(best_config["Inter-parallelism Threads"])), \
               num_gpus, \
               int(float(best_config["Inference Actor Memory Limit (MB)"])), \
               int(float(station_count))
    
