import argparse
import os
from datetime import datetime
import time
import json
from itertools import product

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt    

from sklearn.metrics.cluster import adjusted_rand_score
import sdoclust as sdo

import warnings
warnings.simplefilter(action="ignore", category=FutureWarning)

try:
    from tqdm.auto import tqdm
except ImportError:
    def tqdm(x, **kwargs):
        return x


SDO_FIXED_DEFAULTS = {
    "x": 5, "qv": 0.3, "zeta": 0.6,
    "chi_min": 8, "chi_prop": 0.05, "e": 3, "xc": 5 }

# seeds to repeat runs and estimate stability
DEFAULT_SEEDS = [0, 11, 20]

SDO_PARAM_KEYS = ["k", "x", "qv", "zeta", "chi", "chi_min", "chi_prop", "e", "xc"]

SDO_PARAM_ALIASES = {
    "k":        ["k", "k_", "k_used", "k_final", "k_est", "k_auto"],
    "x":        ["x", "x_", "x_used", "x_final"],
    "qv":       ["qv", "qv_", "qv_used", "qv_final"],
    "zeta":     ["zeta", "zeta_", "zeta_used", "zeta_final"],
    "chi":      ["chi", "chi_", "chi_used", "chi_final", "chi_est"],
    "chi_min":  ["chi_min", "chi_min_", "chiMin", "chi_min_used"],
    "chi_prop": ["chi_prop", "chi_prop_", "chiProp", "chi_prop_used"],
    "e":        ["e", "e_", "epsilon", "epsilon_", "eps", "eps_"],
    "xc":       ["xc", "xc_", "xc_used", "xc_final"] }

BASELINE_CFG = {
    "n_samples":        3000,
    "n_features":       10,
    "n_clusters":       5,
    "outlier_fraction": 0.01,
    "inl_spread":       0.05,
    "outl_spread":      0.5,
    "imbalance_alpha":  0.6,
    "density_ratio":    1.3,
    "collective_frac":  0.1 }

        
def _parse_grid_list(s, cast=float):
    """Parse a comma-separated list like '50,100,200' into typed values."""
    if s is None:
        return []
    items = []
    for raw in str(s).split(","):
        t = raw.strip()
        if not t:
            continue
        try:
            items.append(cast(t))
        except Exception as e:
            raise ValueError(f"could not parse grid value '{raw}' from '{s}': {e}")
    return items

def _parse_seed_list(s):
    """Accepts: '0,1,2' OR [0,1,2] OR '[0,1,2]'."""
    if not s:
        return []
    if isinstance(s, list):
        return [int(x) for x in s]
    if isinstance(s, str):
        s = s.strip()
        if s.startswith("[") and s.endswith("]"):
            s = s[1:-1]
        if not s:
            return []
        return [int(x.strip()) for x in s.split(",") if x.strip()]
    return []

def _safe_tag(v):
    """Create a filesystem-safe string from a parameter value (e.g. 0.6 → '0p6')."""
    if isinstance(v, float):
        return f"{v:g}".replace(".", "p")
    if isinstance(v, (int, np.integer)):
        return str(int(v))
    return str(v)


def build_sdo_grid(k_grid, chi_grid, e_grid, zeta_grid):
    """
    Return a dict of SDOclust configurations for every combination of
    (k, chi, e, zeta), with all other parameters fixed to SDO_FIXED_DEFAULTS.
    Keys are human-readable names like 'SDO_G_k200_chi8_e3_z0p6'.
    """
    algos = {}

    for k, chi, e, zeta in product(k_grid, chi_grid, e_grid, zeta_grid):
        params = dict(SDO_FIXED_DEFAULTS)
        params["k"] = None if k is None else int(k)
        params["chi"] = int(chi)
        params["e"] = float(e) if not float(e).is_integer() else int(e)
        params["zeta"] = float(zeta)

        name = f"SDO_G_k{_safe_tag(k)}_chi{_safe_tag(chi)}_e{_safe_tag(e)}_z{_safe_tag(zeta)}"
        if name in algos:
            raise RuntimeError(f"duplicate algorithm name generated: {name}")
        algos[name] = {"params": params}

    return algos

def build_algorithms(args):
    s = (args.get("sdo_grid") or {})
    k_grid    = s.get("k_grid",    _parse_grid_list("50,100,200", int))
    chi_grid  = s.get("chi_grid",  _parse_grid_list("5,8,12", int))
    e_grid    = s.get("e_grid",    _parse_grid_list("3", float))
    zeta_grid = s.get("zeta_grid", _parse_grid_list("0.6", float))

    if not (k_grid and chi_grid and e_grid and zeta_grid):
        raise ValueError("Invalid SDO grid config")

    return build_sdo_grid(k_grid, chi_grid, e_grid, zeta_grid)

def smooth_fold(X, a=1.5, b=0.75):
    X[:, 1] += a * np.sin(X[:, 0])
    X[:, 0] += b * np.sin(X[:, 1])
    return X

def smooth_holes(X, k=3, strength=0.4):
    r = np.linalg.norm(X, axis=1) / np.sqrt(X.shape[1])  
    X[:, 0] += strength * np.cos(k * r)
    X[:, 1] += strength * np.sin(k * r)
    return X

def generate_data( n_samples, n_features, n_clusters, outlier_fraction,
    inl_spread=0.3, outl_spread=5.0, imbalance_alpha=0.5, density_ratio=2.0, collective_frac=0.2, random_state=None):
    """
    Generate a synthetic dataset with Gaussian clusters and uniform outliers.

    imbalance_alpha : Dirichlet concentration for cluster sizes.
                      Small values (e.g. 0.3) produce strong imbalance,
                      large values (e.g. 10) produce near-equal sizes.
                      None or <= 0 falls back to equal sizes.
    density_ratio   : ratio between the largest and smallest per-cluster sigma.
                      1.0 = all clusters have the same spread (inl_spread).
                      3.0 = spreads range from inl_spread/sqrt(ratio) to
                            inl_spread*sqrt(ratio).
    collective_frac : fraction of outliers generated as small dense groups
                      instead of uniform noise. 0.0 = all outliers are uniform.
    """
    rng = np.random.default_rng(random_state)

    n_inliers  = int(n_samples * (1 - outlier_fraction))
    n_outliers = n_samples - n_inliers
   
    space_scale = n_clusters ** (1 / n_features)
    f = 1+np.log10(n_clusters)
    centers = rng.uniform(-1 * f * space_scale, 1 * f * space_scale, size=(n_clusters, n_features))
   
    # --- cluster sizes (imbalance) ---
    if imbalance_alpha is None or imbalance_alpha <= 0:
        counts = np.full(n_clusters, n_inliers // n_clusters)
        counts[: n_inliers % n_clusters] += 1
    else:
        weights = rng.dirichlet(np.ones(n_clusters) * imbalance_alpha)
        counts  = np.floor(weights * n_inliers).astype(int)
        counts[-1] += n_inliers - counts.sum()   # fix rounding

    # --- per-cluster spreads (variable density) ---
    if density_ratio <= 1.0:
        sigmas = np.full(n_clusters, inl_spread)
    else:
        # log-uniform between inl_spread/sqrt(ratio) and inl_spread*sqrt(ratio)
        # so that the geometric mean stays at inl_spread
        log_half = 0.5 * np.log(density_ratio)
        log_sigmas = rng.uniform(-log_half, log_half, size=n_clusters)
        sigmas = inl_spread * np.exp(log_sigmas)
        
    # --- generate inlier clusters ---
    X_in, y_in = [], []
    nonconvex = True
    for idx, (center, n, sigma) in enumerate(zip(centers, counts, sigmas)):
        if n > 0:
            cluster = rng.normal( loc=center, scale=sigma, size=(n, n_features) ).astype(np.float32)
            a = rng.uniform(1.0, 2.0)
            b = rng.uniform(0.1, 0.8)
            k = rng.integers(1, 4)
            strength = rng.uniform(0, 0.4)
            if nonconvex:
                if n_features >= 2:
                    sub = cluster[:, :2]
                    sub = smooth_fold(sub, a=a, b=b)
                    sub = smooth_holes(sub, k=k, strength=strength)
                    cluster[:, :2] = sub
                else:
                    cluster = smooth_fold(cluster, a=a, b=b)
                    cluster = smooth_holes(cluster, k=k, strength=strength)
    
            if n_features <= 50:
                A = rng.normal(size=(n_features, n_features))
                Q, _ = np.linalg.qr(A)
            else:
                Q = np.eye(n_features)
                
            scales = rng.uniform(0.5, 2.0, size=n_features)
            S = np.diag(scales)
            cluster = cluster @ S @ Q
            
            X_in.append(cluster)
            y_in.extend([idx] * n)

    X_in = np.vstack(X_in)
    y_in = np.array(y_in)

    # --- generate outliers (mix of uniform + collective) ---
    lo = X_in.min() - outl_spread
    hi = X_in.max() + outl_spread

    if n_outliers == 0:
        X_out = np.empty((0, n_features))
    else:
        n_collective = int(n_outliers * collective_frac)
        n_uniform    = n_outliers - n_collective

        parts = []

        if n_uniform > 0:
            parts.append(rng.uniform(lo, hi, size=(n_uniform, n_features)).astype(np.float32))

        if n_collective > 0:
            # place small dense groups at random positions in the same box
            group_size = max(3, n_collective // max(1, n_clusters))
            n_created  = 0
            while n_created < n_collective:
                g       = min(group_size, n_collective - n_created)
                center  = rng.uniform(lo, hi, size=n_features)
                # tight spread: 10% of inl_spread so groups are clearly local
                cluster = rng.normal(center, inl_spread * 0.1, size=(g, n_features)).astype(np.float32)
                parts.append(cluster)
                n_created += g

        X_out = np.vstack(parts)
    
    # --- combine, label, shuffle ---
    X     = np.vstack([X_in, X_out])
    y     = np.concatenate([y_in, np.full(len(X_out), -1, dtype=int)])
    y_bin = (y == -1).astype(int)

    idx        = rng.permutation(len(X))
    X, y, y_bin = X[idx].astype(np.float32), y[idx], y_bin[idx]
    return X, y, y_bin


def _linspace_int(a, b, n):
    """n evenly-spaced integers from a to b, guaranteed strictly increasing."""
    if n == 1:
        return [int(a)]
    vals = np.round(np.linspace(a, b, n)).astype(int).tolist()
    vals[0], vals[-1] = int(a), int(b)
    out = []
    for v in vals:
        out.append(max(v, out[-1] + 1) if out else v)
    return out

def _logspace_int(a, b, n):
    """n integers spaced uniformly in log-space between a and b."""
    if n == 1:
        return [int(a)]

    vals = np.logspace(np.log10(a), np.log10(b), n)
    vals = np.round(vals).astype(int)

    # enforce strictly increasing (important for duplicates after rounding)
    out = []
    for v in vals:
        if not out:
            out.append(int(v))
        else:
            out.append(max(int(v), out[-1] + 1))
    return out

def _linspace_float(a, b, n, decimals=4):
    """n evenly-spaced floats from a to b."""
    if n == 1:
        return [round(float(a), decimals)]
    vals = np.linspace(a, b, n)
    vals[0], vals[-1] = a, b
    return [round(float(v), decimals) for v in vals]

def _logspace_float(a, b, n, decimals=4):
    """n floats spaced uniformly in log-space between a and b."""
    if n == 1:
        return [round(float(a), decimals)]

    vals = np.logspace(np.log10(a), np.log10(b), n)
    return [round(float(v), decimals) for v in vals]


def generate_dataset_configurations(option, i):
    """
    Build a list of dataset configs varying a single factor while keeping all
    others at BASELINE_CFG. Each entry is a complete kwarg dict for generate_data.
    """
    if i < 1:
        raise ValueError("i must be >= 1")

    sweep = {
        "size":      ("n_samples",        lambda: _logspace_int(1000, 1000000, i)),
        "dims":      ("n_features",        lambda: _logspace_int(2, 500, i)),
        "clus":      ("n_clusters",        lambda: _logspace_int(3, 100, i)),
        "outs":      ("outlier_fraction",  lambda: _logspace_float(0.01, 0.40, i)),
        "inspread":  ("inl_spread",        lambda: _linspace_float(0.01, 0.40, i)),
        "outspread": ("outl_spread",       lambda: _linspace_float(0.5,  3.0,  i)),
        "factorial": None,  # handled separately
    }
    
    if option not in sweep:
        raise ValueError(f"unknown analysis_type '{option}'. Choose from: {list(sweep)}")

    cfgs = []
    if option == "factorial":
        for n_clus in _logspace_int(10, 100, i):
            for spread in _linspace_float(0.1, 0.40, i):
                c = dict(BASELINE_CFG)
                c["n_clusters"] = n_clus
                c["inl_spread"] = spread
                cfgs.append(c)
        return cfgs

    else:
        key, grid_fn = sweep[option]
        for v in grid_fn():
            c = dict(BASELINE_CFG)
            c[key] = v
            cfgs.append(c)
        return cfgs


def create_run_directory(base):
    """Create a timestamped sub-folder inside *base* and return its path."""
    os.makedirs(base, exist_ok=True)
    folder = os.path.join(base, f"run{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    os.makedirs(folder, exist_ok=True)
    return folder

def _normalize(x):
    """Convert numpy scalars to native Python types for clean CSV output."""
    if isinstance(x, (np.integer, np.floating, np.bool_)):
        return x.item()
    return x


def extract_effective_params(model, passed):
    """
    Read actual parameter values from a fitted SDOclust model.
    Falls back to *passed* values when model attributes are not found.
    Returns (effective_params, trace) where trace records the source of each value.
    """
    eff, trace = {}, {}

    for key, aliases in SDO_PARAM_ALIASES.items():
        val, src = None, "not_found"
        for attr in aliases:
            if hasattr(model, attr):
                val = _normalize(getattr(model, attr))
                src = f"attr:{attr}"
                break

        if key in passed and passed[key] is not None:
            val = _normalize(passed[key])
            src = f"passed:{key}"
        eff[key]   = val
        trace[key] = src

    return eff, trace


def make_row(cfg, analysis_type, dataset_idx, seed, algo_name, ari, elapsed, eff, passed, error):
    row = dict(cfg)
    row.update({
        "analysis_type": analysis_type,
        "dataset_idx":   int(dataset_idx),
        "seed":          int(seed),
        "algo":          algo_name,
        "ARI":           float(ari),
        "time":          float(elapsed),
        "error":         error })
        
    for p in SDO_PARAM_KEYS:
        row[f"{p}_used"] = eff.get(p)
    for k, v in passed.items():
        if isinstance(v, (int, float, bool, str)) or v is None:
            row[f"passed_{k}"] = _normalize(v)
    return row

def save_dataset_scatter(X, y_bin, out_path):
    if X.shape[1] < 2:
        return
    plt.figure(figsize=(4, 4))
    colors = np.where(y_bin == 1, "#d62728", "#1f77b4")
    plt.scatter(X[:, 0], X[:, 1], c=colors, s=10, alpha=0.8, linewidths=0)
    plt.xticks([])
    plt.yticks([])
    plt.tight_layout()
    plt.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close()

def run_one_analysis(analysis_type, args, algos):
    """
    Run the full sensitivity experiment for one analysis type.
    Returns a long-format DataFrame with one row per (seed, dataset, algorithm).
    """
    configs = generate_dataset_configurations(analysis_type, int(args["i"]))
    seeds   = list(args["seed_list"])
    rows = []

    for seed in tqdm(seeds, desc=f"{analysis_type}: seeds", unit="seed", dynamic_ncols=True):
        for j, cfg in enumerate(tqdm(configs, desc=f"  datasets (seed={seed})", unit="ds", leave=False, dynamic_ncols=True) ):
        
            X,y,y_bin = generate_data(**cfg, random_state=seed+j)
            
            # control plot
            control_plot = False
            if control_plot:
                if j % int(args["i"]/3) == 0:
                    img_dir = os.path.join(args["results_folder"], "datasets_png", analysis_type)
                    os.makedirs(img_dir, exist_ok=True)
                    img_path = os.path.join( img_dir, f"ds_seed{seed}_i{j}.png")
                    save_dataset_scatter(X, y_bin, img_path)

            for name, algo_cfg in tqdm( algos.items(), desc="    algos", unit="algo", leave=False, dynamic_ncols=True ):
                params = dict(algo_cfg["params"])
                # clamp k so it never exceeds the dataset size
                if "k" in params and params["k"] is not None:
                    params["k"] = min(int(params["k"]), max(5, X.shape[0] - 1))

                model = sdo.SDOclust(**params)
                error = None
                t0    = time.perf_counter()
                try:
                    preds = model.fit_predict(X)
                except Exception as exc:
                    error = str(exc)[:200]
                    preds = np.full(X.shape[0], -1, dtype=int)
                elapsed = time.perf_counter() - t0

                # compute ARI on inliers only
                mask = y_bin == 0
                try:
                    ari = adjusted_rand_score(y[mask], preds[mask])
                except Exception:
                    ari = 0.0

                eff, _ = extract_effective_params(model, passed=params)
                rows.append(make_row(cfg, analysis_type, j, seed, name, ari, elapsed, eff, params, error))
                del model

    return pd.DataFrame(rows)


def parse_args():
    parser = argparse.ArgumentParser(description="SDOclust sensitivity analysis.")
    parser.add_argument("--config", required=True)
    
    with open(parser.parse_args().config) as f:
        cfg = json.load(f)

    if "analysis_type" not in cfg:
        raise ValueError("analysis_type missing in config.json")

    cfg["i"] = cfg.get("i", 10)
    cfg["o"] = cfg.get("o", "results")
    cfg["seed_list"] = _parse_seed_list(cfg.get("seeds", "0,1,2,3,4")) or list(DEFAULT_SEEDS)
    cfg["results_folder"] = create_run_directory(cfg["o"])
    cfg.setdefault("sdo_grid", None)

    print(f"\nanalysis={cfg['analysis_type']} | i={cfg['i']} | seeds={len(cfg['seed_list'])} | out={cfg['results_folder']}\n")

    with open(os.path.join(cfg["results_folder"], "args.json"), "w") as f:
        json.dump(cfg, f, indent=2)

    return cfg
    

def main():
    args = parse_args()

    all_types = ["size", "dims", "clus", "outs", "inspread", "outspread", "factorial"]
    selected  = all_types if args["analysis_type"] == "all" else [args["analysis_type"]]
    if args["analysis_type"] not in all_types + ["all"]:
        raise ValueError(f"analysis_type must be one of {all_types} or 'all'")

    algos = build_algorithms(args)
    print(f"SDO grid: {len(algos)} combinations")
    for name, cfg in sorted(algos.items()):
        p = cfg["params"]
        print(f"  {name}: k={p['k']}, chi={p['chi']}, e={p['e']}, zeta={p['zeta']}")
    
    out_csv = os.path.join(args["results_folder"], "ALL_combinations_LONG.csv")
    first = True
    for at in selected:
        df = run_one_analysis(at, args, algos)
        df.to_csv( out_csv, mode="w" if first else "a",  header=first, index=False )
        first = False
        del df       
        print(f"\nsaved: {out_csv}")


if __name__ == "__main__":
    main()