"""
pycutwm.core
------------

Main routines for managing grid definition, compiling CUDA code, and parameter sweeps
for the cuTWM(3+1)D simulation package.

Functions:
    - write_grid_on_file: Write grid/cuda header from JSON config
    - compile_cutwm: Compile the cuTWM CUDA/C++ code
    - run_cutwm: Run a single simulation
    - run_cutwm_sweep: Parameter sweep manager

Dependencies:
    - Python >= 3.7
    - numpy
    - nvidia-smi in PATH (for compile_cutwm)
    - Custom utils (pycutwm.utils.load_config, save_config)

Author: Alfredo Sánchez
Date: 16/10/2025
License: MIT
"""


import os
import shutil
import json
import subprocess
import time
import glob
import re
import sys
import numpy as np

from .utils import load_config, save_config

# ---------------- Logging: real-time capture of stdout/stderr from a command ----------------
def run_and_tee(cmd, log_path, cwd=None, env=None):
    """
    Runs `cmd`, capturing everything printed to the terminal and saving it to `log_path`,
    while also displaying it live on the screen.
    """
    log_dir = os.path.dirname(log_path)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)

    with open(log_path, "a", encoding="utf-8", buffering=1) as lf:
        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            cwd=cwd,
            env=env
        )
        # Read line by line and output to console and file
        for line in proc.stdout:
            sys.stdout.write(line)
            lf.write(line)
        returncode = proc.wait()

    return returncode


def run_cutwm(config_path="./config.json", exe_path="./cuTWM", log_path=None, cwd=None, env=None):
    """
    Run the main executable with config.json.

    If `log_path` is None, it is automatically generated next to the config:
        <same_folder>/<same_base>.log
    """
    os.environ["LD_LIBRARY_PATH"] = os.environ["CONDA_PREFIX"] + "/lib:" + os.environ.get("LD_LIBRARY_PATH", "")

    # Derive default log name next to the JSON config
    if log_path is None:
        cfg_abs = os.path.abspath(config_path)
        base, _ = os.path.splitext(cfg_abs)
        log_path = base + ".log"

    cmd = [exe_path, config_path]
    print("Running:", " ".join(cmd))
    rc = run_and_tee(cmd, log_path=log_path, cwd=cwd, env=env)
    if rc != 0:
        raise RuntimeError(f"Simulation failed (exit code {rc}). Log: {log_path}")
    print(f"\n\nSimulation finished. Log saved to: {log_path}\n\n")


def move_h5_files(dest_folder):
    """
    Move all .h5 files from current directory to dest_folder.
    """
    # Create the folder if it does not exist
    if not os.path.exists(dest_folder):
        os.makedirs(dest_folder)

    # Move files from the current folder
    for file in os.listdir("."):
        if file.endswith(".h5"):
            src = os.path.join(".", file)
            dst = os.path.join(dest_folder, file)
            shutil.move(src, dst)
            print(f"Moved: {file} -> {dest_folder}")

# ---------- Helpers ----------
def _format_value(v, decimals=1) -> str:
    """
    Format the value for filenames. For floats, set the number of decimals.
    Example: 1   -> "1"
             1.0 -> "1.0" (if decimals=1)
             5.25 -> "5.3" (if decimals=1)
    """
    if isinstance(v, float):
        return f"{v:.{decimals}f}"
    return str(v)


def _sanitize_value_str(s: str) -> str:
    """
    Allows letters, numbers, period, hyphen, underscore, + and =.
    Does NOT replace the period.
    """
    return re.sub(r"[^A-Za-z0-9.\-+=_]", "_", s)


def _value_str(val, decimals=1):
    if isinstance(val, (list, tuple, np.ndarray)):
        return "__".join([f"{v:.{decimals}f}" for v in val])
    else:
        return f"{val:.{decimals}f}"

# ---------- Sweep ----------
def run_cutwm_sweep(config_path, param_path, values,
                    exe_path="path",
                    out_dir_base="./results/sweep",
                    keep_config=True,
                    source_dir=".",              # where to collect .dat files if exe does not respect output_directory
                    patterns=("*.dat",),         # which file patterns to move after run
                    mtime_tolerance_s=1.0,       # timestamp tolerance for file modification
                    value_decimals=1             # number of decimals for value names
                    ):
    """
    Sweep a single parameter and leave the output of each run in:
        <out_dir_base>/simulation_<leaf>_<value>/

    - If the executable respects cfg['simulation']['output_directory'], the .dat
      files will appear directly in that folder.
    - If NOT (writes elsewhere, e.g. CWD), try to move .dat files created/modified from `source_dir`.

    Also generates a .log with the SAME name as the temporary JSON used for the run,
    saving it in the same output folder (out_dir).
    """
    cfg = load_config(config_path)

    keys = param_path.split("/")
    leaf = keys[-1]

    os.makedirs(out_dir_base, exist_ok=True)

    for v in values:
        # 1) Apply the value in the config (note: if your executable modifies the JSON in memory,
        #    consider reloading cfg with load_config(config_path) every iteration)
        ptr = cfg
        for k in keys[:-1]:
            ptr = ptr[k]
        ptr[leaf] = v

        # 2) Build the name and folder (keep the period in floats)
        val_str = _value_str(v, decimals=value_decimals)
        folder_name = f"simulation_{leaf}_{val_str}"
        out_dir = os.path.join(out_dir_base, folder_name)
        os.makedirs(out_dir, exist_ok=True)

        # 3) Update simulation metadata
        base_name = cfg.get('simulation', {}).get('name', 'sim')
        sim_name = f"{base_name}_{leaf}_{val_str}"
        cfg['simulation']['name'] = sim_name
        cfg['simulation']['output_directory'] = out_dir

        # 4) Save config in the output folder
        temp_config = os.path.join(out_dir, f"config_{leaf}_{val_str}.json")
        save_config(cfg, temp_config)

        # 5) Run, also registering a log with the SAME base name as the JSON
        #    (e.g., config_pump_power_W_1.0.json -> config_pump_power_W_1.0.log)
        log_path = os.path.splitext(temp_config)[0] + ".log"
        t0 = time.time()
        run_cutwm(config_path=temp_config, exe_path=exe_path, log_path=log_path)

        # 6) If there are no .dat files in out_dir, move newly created ones from source_dir
        # has_h5_in_out = any(fn.lower().endswith(".dat") for fn in os.listdir(out_dir))
        has_h5_in_out = any(fn.lower().endswith(".h5") for fn in os.listdir(out_dir))
        if not has_h5_in_out:
            for pat in patterns:
                for f in glob.glob(os.path.join(source_dir, pat)):
                    try:
                        # Move new/modified files since the start of the run
                        if os.path.getmtime(f) >= (t0 - mtime_tolerance_s):
                            dst = os.path.join(out_dir, os.path.basename(f))
                            if os.path.exists(dst):
                                base, ext = os.path.splitext(dst)
                                dst = f"{base}__{leaf}_{val_str}{ext}"
                            shutil.move(f, dst)
                    except FileNotFoundError:
                        pass
        move_h5_files(out_dir)                
        # 7) Optional config cleanup
        if not keep_config:
            try:
                os.remove(temp_config)
            except OSError:
                pass
