import gc
import numpy as np
import pandas as pd
from tqdm import tqdm
import os, warnings, traceback, argparse, multiprocessing as mp
from utils import generate_dataset_configurations, generate_data, generate_interaction_configurations, run_ad, run_cl, set_analysis_algorithms
import tracemalloc, time
from multiprocessing import shared_memory
from sklearn.preprocessing import RobustScaler

warnings.filterwarnings("ignore", category=FutureWarning)

# Force single-threaded BLAS/MKL so subprocesses don't over-subscribe CPU cores
os.environ["OMP_NUM_THREADS"]     = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"]     = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

N_JOBS   = 1
TIMEOUT  = 5 * 60   # 5 minutes


def run_with_timeout(func, args=(), kwargs={}, timeout=TIMEOUT, seed=None):
    # Move large arrays into shared memory so the subprocess can access them
    # without pickling (avoids expensive serialization for big datasets)
    shms = []
    new_args = []
    for arg in args:
        if isinstance(arg, np.ndarray) and arg.nbytes > 50 * 1024**2:  # >50 MB
            shm = shared_memory.SharedMemory(create=True, size=arg.nbytes)
            shared_arr = np.ndarray(arg.shape, dtype=arg.dtype, buffer=shm.buf)
            np.copyto(shared_arr, arg)
            new_args.append((shm.name, arg.shape, arg.dtype))  # pass a descriptor tuple
            shms.append(shm)
        else:
            new_args.append(arg)

    q = mp.Queue()
    p = mp.Process(target=_wrapper_shm, args=(q, func, new_args, kwargs, seed))
    p.start(); p.join(timeout)

    # Always release shared memory blocks, even on timeout or error
    for shm in shms:
        shm.close(); shm.unlink()

    if p.is_alive():
        p.terminate(); p.join()
        raise TimeoutError("Algorithm exceeded time limit.")
    if q.empty():
        raise RuntimeError("No result returned from subprocess.")

    status, payload = q.get()
    if status == 'success':
        return (*payload[0], payload[1], payload[2])  # (*scores, duration, peak_memory)
    if status == 'memory':  raise MemoryError(payload)
    raise Exception(payload)


def _wrapper_shm(queue, func, args, kwargs, seed=None):
    """Subprocess entry point: reconstruct shared-memory arrays, run func, report back."""
    try:
        if seed is not None:
            np.random.seed(seed)
        real_args = []
        shms = []
        for arg in args:
            # Rebuild arrays passed as shared-memory descriptors.
            if isinstance(arg, tuple) and len(arg) == 3 and isinstance(arg[0], str):
                name, shape, dtype = arg
                shm = shared_memory.SharedMemory(name=name)
                arr = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
                real_args.append(arr)
                shms.append(shm)
            else:
                real_args.append(arg)

        # Measure Python-level peak memory during the algorithm run.
        tracemalloc.start()
        t0 = time.perf_counter()
        result = func(*real_args, **kwargs)
        t1 = time.perf_counter()
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()

        for shm in shms:
            shm.close()

        queue.put(('success', (result, t1 - t0, peak)))
    except MemoryError as e:
        queue.put(('memory', str(e)))
    except Exception:
        queue.put(('error', traceback.format_exc()))


def _set_status(row, name, status, cols, scores, duration):
    """Write metric scores, elapsed time, and status into a result row dict."""
    for col, score in zip(cols, scores):
        row[col] = score
    row[f"{name}_time"]   = duration
    row[f"{name}_status"] = status
    return row

def _upsert_row(df, row, config, rep, output_file):
    """Replace the matching row in df (by config + rep) and persist to CSV."""
    n_samples, n_features, n_clusters, outlier_fraction = config
    mask = ((df['n_samples']        == n_samples)        &
            (df['n_features']       == n_features)       &
            (df['n_clusters']       == n_clusters)        &
            (df['outlier_fraction'] == outlier_fraction) &
            (df['rep']              == rep))

    df = pd.concat([df[~mask], pd.DataFrame([row])], ignore_index=True)
    df.to_csv(output_file, index=False)

    return df

def _run_loop(df, row, config, scenario, rep, algorithms, task, metric_suffixes, run_fn, run_args_fn, output_file, skip_algos, seed):
    for name, cfg in tqdm(algorithms.items(), desc=f"{task}", leave=False):
        skip_key = (config, name)
        
        # Skip algorithms that previously timed out or ran out of memory on this config
        if skip_key in skip_algos:
            status = skip_algos[skip_key]
            print(f"  Skipping {task}: {name} - {config} {scenario} - "  f"(prev {status})")

            cols = [f"{name}_{s}" for s in metric_suffixes]
            row = _set_status(row, name, status, cols, [np.nan] * len(cols), np.nan)
            df = _upsert_row(df, row, config, rep, output_file)
            continue

        
        # Skip if primary metric column already has a value (resume-friendly)
        first_col = f"{name}_{metric_suffixes[0]}"
        if first_col in row and not pd.isna(row[first_col]):
            print(f"  Skipping {task}: {name} - {config} {scenario} - (already computed)")
            continue

        cols = [f"{name}_{s}" for s in metric_suffixes]
            
        try:
            *scores, duration, memory = run_with_timeout(run_fn, args=run_args_fn(name, cfg), seed=seed)
            row = _set_status(row, name, "success", cols, scores, duration)
            row[f"{name}_memory_mb"] = memory / (1024**2)
        except TimeoutError:
            print(f"  Failed {task}: {name} - {config} {scenario} - timeout")
            row = _set_status(row, name, "timeout", cols, [np.nan]*len(cols), np.nan)
            # Skip this algorithm in later repetitions of the same configuration.
            skip_algos[skip_key] = "timeout" 
        except MemoryError:
            print(f"  Failed {task}: {name} - {config} {scenario} - memory")
            row = _set_status(row, name, "memory", cols, [np.nan]*len(cols), np.nan)
            # Skip this algorithm in later repetitions of the same configuration.
            skip_algos[skip_key] = "memory"
        except Exception as e:
            print(f"  Failed {task}: {name} - {config} {scenario} - {e}")
            row = _set_status(row, name, "error", cols, [np.nan]*len(cols), np.nan)

        df = _upsert_row(df, row, config, rep, output_file)
    return df, row


def run_analysis(analysis_type="size", scenario="easy", mode="both", n_reps=5, config_index=None, rep_index=None, output_file=None, random_seed=0):
    """
    Run anomaly detection / clustering experiments on synthetic datasets.

    Args:
        analysis_type: Parameter to vary — 'size' | 'dims' | 'outs' | 'clus'
        scenario:      Dataset difficulty — 'easy' | 'hard'
        mode:          Tasks to run — 'ad' | 'cl' | 'both'
        n_reps:        Number of independent repetitions
        config_index:  If set, run only this configuration (0-based)
        rep_index:     If set, run only this repetition (0-based); requires config_index
        output_file:   CSV path override
    """
    
    assert mode in ("ad", "cl", "both"), "mode must be 'ad', 'cl', or 'both'"
    assert scenario in ("easy", "hard"), "scenario must be 'easy' or 'hard'"

    results_dir = os.path.join("results", f"results_{analysis_type}_{scenario}")
    os.makedirs(results_dir, exist_ok=True)
    output_file = output_file or os.path.join(results_dir, f"results_{analysis_type}_{scenario}.csv")

    if analysis_type == "interaction":
        configs = generate_interaction_configurations(scenario, True)    
    else:
        configs = generate_dataset_configurations(analysis_type, scenario)

    if config_index is not None:
        configs = [configs[config_index]]

    columns=["n_samples", "n_features", "n_clusters", "outlier_fraction", "rep"]

    # Load existing results so interrupted runs can be resumed
    if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
         try:
             df = pd.read_csv(output_file, na_values=["", "nan", "NaN"]).replace("", np.nan)
         except pd.errors.EmptyDataError:
             df = pd.DataFrame(columns=columns)
    else:
        df = pd.DataFrame(columns=columns)
    
    reps = [rep_index] if rep_index is not None else list(range(n_reps))
        
    skip_algos = {}  # key: (config_tuple, algo_name) → failure reason; persists across reps
    
    total = len(configs) * len(reps)
    with tqdm(total=total, desc="Overall progress") as pbar:
    
        for i, config in enumerate(configs):
            n_samples, n_features, n_clusters, outlier_fraction = config
        
            for rep in reps:
                seed = random_seed + rep
                np.random.seed(seed)
                print(f"\n[{analysis_type}|{scenario}] Config {i+1}/{len(configs)}: "
                      f"{n_samples} samples, {n_features} dims, {n_clusters} clusters, "
                      f"{outlier_fraction*100:.0f}% outliers  |  rep {rep+1}/{n_reps} (seed={seed})")

                mask = ((df['n_samples'] == n_samples) & (df['n_features'] == n_features) & (df['n_clusters'] == n_clusters) &
                        (df['outlier_fraction'] == outlier_fraction) & (df['rep'] == rep))
                if not df.empty and mask.any():
                    row = df.loc[mask].iloc[0].to_dict()
                else:
                    row = { "n_samples": n_samples, "n_features": n_features, "n_clusters": n_clusters,
                        "outlier_fraction": outlier_fraction, "rep": rep, "seed": seed }
                        
                _, _, Xraw, yo, yc = generate_data(n_samples, n_features, n_clusters, outlier_fraction, scenario, seed)
                Xo = RobustScaler().fit_transform(Xraw)
                ad_algorithms, clu_algorithms = set_analysis_algorithms(n_clusters, len(yo), dims=n_features, n_jobs=N_JOBS, seed=seed)

                if mode in ("ad", "both"):
                    df, row = _run_loop(
                        df, row, config, scenario, rep, 
                        algorithms=ad_algorithms, task="AD", metric_suffixes=["roc", "ap"],
                        run_fn=run_ad, run_args_fn=lambda n, cfg: (cfg, Xo, yo),  # n unused for AD
                        output_file=output_file, skip_algos=skip_algos, seed=seed )
                if mode in ("cl", "both"):
                    df, row = _run_loop(
                        df, row, config, scenario, rep, 
                        algorithms=clu_algorithms, task="CL", metric_suffixes=["adjRand", "ami"],
                        run_fn=run_cl, run_args_fn=lambda n, cfg: (n, cfg, Xo, yc),
                        output_file=output_file, skip_algos=skip_algos, seed=seed )
                gc.collect()
                pbar.update(1)
            
    print("\nDone!")

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Run anomaly detection and clustering experiments.")
    parser.add_argument("--analysis_type", required=True, choices=["size", "dims", "outs", "clus", "interaction"])
    parser.add_argument("--scenario",      default="easy", choices=["easy", "hard"])
    parser.add_argument("--mode",          default="both", choices=["ad", "cl", "both"])
    parser.add_argument("--n_reps",        type=int, default=5)
    parser.add_argument("--config_index",  type=int, default=None)
    parser.add_argument("--rep_index",     type=int, default=None)
    parser.add_argument("--output_file",   type=str, default=None)
    parser.add_argument("--random_seed", type=int, default=0)
    args = parser.parse_args()

    run_analysis(args.analysis_type, args.scenario, args.mode, args.n_reps, args.config_index, args.rep_index, args.output_file, args.random_seed)