#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Generate 3 datasets and compute reference statistics for each ALS Statistics instrument.
Writes JSON files into ./out/set-001.json, set-002.json, set-003.json

Dependencies:
  - numpy
  - scipy
  - (optional) scikit-learn (for logistic regression)
  - (optional) hdbscan     (for HDBSCAN; if missing, marked unavailable)

Run:
  python make_goldens_all.py
"""
from __future__ import annotations

import json, os, math, argparse, random
from dataclasses import dataclass, asdict
from typing import Dict, Any, List, Tuple, Optional

import numpy as np
from numpy.typing import NDArray
from scipy import stats
import hdbscan

# ---------- helpers ----------

def to_float(x):
    if isinstance(x, (np.floating, np.integer)):
        return float(x)
    return x

def listify(arr: NDArray) -> List[float]:
    return [to_float(v) for v in np.asarray(arr).ravel().tolist()]

def ensure_dir(p: str):
    os.makedirs(p, exist_ok=True)

# ---------- core stats used by several modules ----------

def cronbach_alpha(items: NDArray) -> float:
    """
    items: (n_samples, k_items)
    """
    n, k = items.shape
    if k < 2:
        return float('nan')
    variances = items.var(axis=0, ddof=1)
    total = items.sum(axis=1)
    var_total = total.var(ddof=1)
    if var_total == 0:
        return float('nan')
    return (k / (k - 1.0)) * (1.0 - (variances.sum() / var_total))

def cronbach_if_deleted(items: NDArray, names: List[str]) -> Dict[str, float]:
    out = {}
    for i, name in enumerate(names):
        sub = np.delete(items, i, axis=1)
        out[name] = float(cronbach_alpha(sub))
    return out

def welch_anova(*groups: NDArray) -> Tuple[float, float, float, float]:
    """
    Returns F, df1, df2, p for Welch one-way ANOVA.
    """
    k = len(groups)
    if k < 2:
        return (float('nan'), float('nan'), float('nan'), float('nan'))
    means = np.array([np.mean(g) for g in groups], dtype=float)
    ns = np.array([len(g) for g in groups], dtype=float)
    variances = np.array([np.var(g, ddof=1) for g in groups], dtype=float)
    wi = ns / variances
    W = wi.sum()
    ybar = (wi * means).sum() / W
    numerator = np.sum(wi * (means - ybar) ** 2) / (k - 1.0)
    u = np.sum((1.0 / (ns - 1.0)) * (1.0 - (wi / W)) ** 2)
    denom = 1.0 + (2.0 * (k - 2.0) / (k ** 2 - 1.0)) * u
    F = numerator / denom
    df1 = k - 1.0
    df2 = (k ** 2 - 1.0) / (3.0 * u) if u > 0 else float('inf')
    p = 1.0 - stats.f.cdf(F, df1, df2)
    return float(F), float(df1), float(df2), float(p)

def mad(x: NDArray) -> float:
    # median absolute deviation (about the median), unscaled
    x = np.asarray(x, dtype=float)
    med = np.median(x)
    return float(np.median(np.abs(x - med)))

def mad_distance(a: NDArray, b: NDArray, eps: float = 1e-12) -> float:
    """
    Heuristic distance between two equal-length series based on median absolute deviations.
    d = median(|a-b|) / (MAD(a) + MAD(b) + eps)
    """
    a = np.asarray(a, dtype=float)
    b = np.asarray(b, dtype=float)
    num = np.median(np.abs(a - b))
    denom = mad(a) + mad(b) + eps
    return float(num / denom)

def pairwise_mad_distances(cols: Dict[str, NDArray]) -> Tuple[List[str], NDArray]:
    names = list(cols.keys())
    m = len(names)
    D = np.zeros((m, m), dtype=float)
    for i in range(m):
        for j in range(i + 1, m):
            d = mad_distance(cols[names[i]], cols[names[j]])
            D[i, j] = D[j, i] = d
    return names, D

# Simple DBSCAN on precomputed distance matrix (columns-as-points)
def dbscan_from_distance(D: NDArray, eps: float = 0.5, min_pts: int = 2) -> List[int]:
    """
    Very small DBSCAN implementation over a symmetric distance matrix.
    Returns labels: -1 noise, 1..k cluster ids.
    """
    n = D.shape[0]
    labels = [0] * n   # 0=unvisited, -1=noise, 1..C clusters
    cluster_id = 0

    def neighbors(i: int) -> List[int]:
        return [j for j in range(n) if j != i and D[i, j] <= eps]

    for i in range(n):
        if labels[i] != 0:
            continue
        nbrs = neighbors(i)
        if len(nbrs) + 1 < min_pts:
            labels[i] = -1
            continue
        cluster_id += 1
        labels[i] = cluster_id
        seed = set(nbrs)
        while seed:
            j = seed.pop()
            if labels[j] == -1:
                labels[j] = cluster_id
            if labels[j] != 0:
                continue
            labels[j] = cluster_id
            nbrs2 = neighbors(j)
            if len(nbrs2) + 1 >= min_pts:
                for q in nbrs2:
                    if labels[q] == 0:
                        seed.add(q)
    return labels

# Optional logistic regression (sklearn), HDBSCAN
def try_logistic(X: NDArray, y: NDArray) -> Dict[str, Any]:
    try:
        from sklearn.linear_model import LogisticRegression
        model = LogisticRegression(max_iter=1000)
        model.fit(X, y)
        acc = float(model.score(X, y))
        return {
            "available": True,
            "coefficients": listify(np.concatenate(([model.intercept_[0]], model.coef_[0]))),
            "accuracy": acc
        }
    except Exception as e:
        return {"available": False, "reason": str(e)}

# def try_hdbscan(dist_names: List[str], D: NDArray, min_cluster_size: int = 2) -> Dict[str, Any]:
#     try:
#         import hdbscan
#         # Convert distance matrix to condensed distance for hierarchy? hdbscan expects points or condensed;
#         # We'll use precomputed distances by embedding with MDS-like hack is overkill.
#         # Instead, we can feed as if it's a feature space via classical MDS—skip; mark unavailable.
#         return {"available": False, "reason": "Precomputed distance not directly supported in this script"}
#     except Exception as e:
#         return {"available": False, "reason": str(e)}


def try_hdbscan(dist_names: List[str], D: NDArray, min_cluster_size: int = 2) -> Dict[str, Any]:
    """
    Run HDBSCAN on a **precomputed** distance matrix.

    Parameters
    ----------
    dist_names : list[str]
        Names of series/columns in the same order as the distance matrix rows/cols.
    D : np.ndarray, shape (n, n)
        Square, symmetric, non-negative distance matrix (zeros on the diagonal).
    min_cluster_size : int
        HDBSCAN min_cluster_size parameter (default: 2).

    Returns
    -------
    dict
        {
          "available": bool,
          "reason": str (if available=False),
          "version": str | None,
          "min_cluster_size": int,
          "n_clusters": int,
          "labels": { name: int },           # -1 = noise, >=1 cluster ids (library uses 0.. but we reindex to 1..)
          "probabilities": { name: float },  # membership strength [0,1]
          "noise": [name, ...],
          "clusters": [ { "id": int, "columns": [name,...] }, ... ]
        }
    """
    try:
        # ---- validate matrix shape ----
        D = np.asarray(D, dtype=float)
        n = D.shape[0]
        if D.ndim != 2 or D.shape[0] != D.shape[1]:
            return {"available": False, "reason": "D must be a square (n×n) distance matrix"}
        if len(dist_names) != n:
            return {"available": False, "reason": "len(dist_names) must match D.shape[0]"}
        if np.any(np.isnan(D)) or np.any(np.isinf(D)):
            return {"available": False, "reason": "Distance matrix contains NaN/Inf"}
        if np.any(D < 0):
            return {"available": False, "reason": "Distance matrix must be non-negative"}
        # enforce symmetry & zero diagonal (tolerant)
        if not np.allclose(D, D.T, atol=1e-12, rtol=1e-12):
            return {"available": False, "reason": "Distance matrix is not symmetric within tolerance"}
        np.fill_diagonal(D, 0.0)

        # ---- fit HDBSCAN with precomputed distances ----
        clusterer = hdbscan.HDBSCAN(
            metric="precomputed",
            min_cluster_size=int(min_cluster_size)
        ).fit(D)

        raw_labels = clusterer.labels_            # -1 for noise, 0..C-1 for clusters
        probs = getattr(clusterer, "probabilities_", None)
        if probs is None:
            probs = np.ones_like(raw_labels, dtype=float)

        # Reindex clusters from 0.. to 1.. (to match DBSCAN-style positive ids)
        reindex = {}
        next_id = 1
        for lab in raw_labels:
            if lab >= 0 and lab not in reindex:
                reindex[lab] = next_id
                next_id += 1

        mapped_labels = np.array([(-1 if lab < 0 else reindex[lab]) for lab in raw_labels], dtype=int)

        # Build outputs
        labels_by_name = { name: int(mapped_labels[i]) for i, name in enumerate(dist_names) }
        probs_by_name  = { name: float(probs[i]) for i, name in enumerate(dist_names) }

        # Clusters and noise lists
        clusters_map: Dict[int, List[str]] = {}
        noise: List[str] = []
        for i, name in enumerate(dist_names):
            lab = int(mapped_labels[i])
            if lab == -1:
                noise.append(name)
            else:
                clusters_map.setdefault(lab, []).append(name)

        clusters_list = [
            {"id": cid, "columns": sorted(cols)}
            for cid, cols in sorted(clusters_map.items(), key=lambda kv: kv[0])
        ]

        return {
            "available": True,
            "reason": "",
            "version": getattr(hdbscan, "__version__", None),
            "min_cluster_size": int(min_cluster_size),
            "n_clusters": len(clusters_list),
            "labels": labels_by_name,
            "probabilities": probs_by_name,
            "noise": sorted(noise),
            "clusters": clusters_list,
        }

    except ImportError as e:
        return {"available": False, "reason": f"hdbscan not installed: {e}"}
    except Exception as e:
        return {"available": False, "reason": f"failed to run hdbscan: {e}"}


# ---------- datasets ----------

def gaussian(rng: random.Random, mean=0.0, sd=1.0) -> float:
    # Box-Muller
    u1 = rng.random() or 1e-12
    u2 = rng.random() or 1e-12
    z = math.sqrt(-2.0 * math.log(u1)) * math.cos(2.0 * math.pi * u2)
    return mean + sd * z

def make_dataset(kind: int, rows: int, seed: int) -> Dict[str, Any]:
    """
    kind:
      1 - correlated normals + 3 groups (unequal variances)
      2 - heavier tails, stronger group shifts, before/after larger delta
      3 - mixed discrete/continuous with outliers sprinkled
    """
    rng = random.Random(seed)

    X, Y, group, score, before, after = [], [], [], [], [], []
    Q1, Q2, Q3, Q4 = [], [], [], []
    for i in range(rows):
        if kind == 1:
            # correlated X,Y
            x = gaussian(rng, 50, 10)
            noise = gaussian(rng, 0, 12 * math.sqrt(1 - 0.65**2))
            y = 30 + (0.65 * (12/10)) * (x - 50) + noise
            g = rng.randrange(3)
            base = 70 + gaussian(rng, 0, 8 * (1.0 + 0.4*g))   # unequal var by group
            shift = {0: 0, 1: +6, 2: -4}[g]
            b = 60 + gaussian(rng, 0, 7)
            d = 4 + gaussian(rng, 0, 2)
        elif kind == 2:
            x = gaussian(rng, 40, 15)
            y = 5 + 1.1 * x + gaussian(rng, 0, 18)
            g = rng.randrange(3)
            base = 65 + gaussian(rng, 0, 10 * (1.0 + 0.6*g))
            shift = {0: -2, 1: +10, 2: +3}[g]
            b = 55 + gaussian(rng, 0, 10)
            d = 8 + gaussian(rng, 0, 3)
        else:
            # kind 3: mix, add occasional outliers
            x = 30 + (rng.random() * 40)
            y = 10 + 0.8 * x + gaussian(rng, 0, 6)
            if rng.random() < 0.03:
                y += (25 if rng.random() < 0.5 else -20)
            g = rng.randrange(3)
            base = 68 + gaussian(rng, 0, 7 * (1.0 + 0.5*g))
            shift = {0: 0, 1: +3, 2: +7}[g]
            b = 58 + gaussian(rng, 0, 6)
            d = 2 + gaussian(rng, 0, 2)

        X.append(x)
        Y.append(y)
        group.append(g)
        score.append(base + shift)
        before.append(b)
        after.append(b + d)

        # reliability items from latent
        z = gaussian(rng, 0, 1)
        Q1.append(10 + 2.0*z + gaussian(rng,0,0.8))
        Q2.append( 9 + 1.8*z + gaussian(rng,0,0.9))
        Q3.append(11 + 2.1*z + gaussian(rng,0,0.7))
        Q4.append( 8 + 1.7*z + gaussian(rng,0,1.0))

    return {
        "meta": {"kind": kind, "rows": rows, "seed": seed},
        "data": {
            "X": X, "Y": Y, "group": group, "score": score,
            "before": before, "after": after,
            "Q1": Q1, "Q2": Q2, "Q3": Q3, "Q4": Q4
        }
    }

# ---------- compute per instrument ----------

def descriptive_block(arr: NDArray) -> Dict[str, Any]:
    arr = np.asarray(arr, dtype=float)
    n = arr.size
    mean = float(np.mean(arr))
    med = float(np.median(arr))
    var_pop = float(np.var(arr, ddof=0))
    var_sam = float(np.var(arr, ddof=1)) if n > 1 else float('nan')
    sd_pop = float(np.std(arr, ddof=0))
    sd_sam = float(np.std(arr, ddof=1)) if n > 1 else float('nan')
    q1 = float(np.percentile(arr, 25))
    q3 = float(np.percentile(arr, 75))
    p10 = float(np.percentile(arr, 10))
    p90 = float(np.percentile(arr, 90))
    iqr = float(stats.iqr(arr))
    mad_v = float(stats.median_abs_deviation(arr, scale=1.0, nan_policy='omit'))
    skew = float(stats.skew(arr, bias=False)) if n > 2 else float('nan')
    kurt = float(stats.kurtosis(arr, fisher=True, bias=False)) if n > 3 else float('nan')
    gmean = float(stats.gmean(np.clip(arr, 1e-12, None)))  # avoid zeros/negatives
    try:
        hmean = float(stats.hmean(np.clip(arr, 1e-12, None)))
    except Exception:
        hmean = float('nan')
    zscores = listify((arr - mean) / (sd_sam if sd_sam not in (0.0, float('nan')) else (sd_pop or 1e-12)))
    outliers_z3_idx = [int(i) for i, z in enumerate(zscores) if abs(z) > 3]
    lo, hi = q1 - 1.5 * iqr, q3 + 1.5 * iqr
    outliers_iqr_idx = [int(i) for i, v in enumerate(arr) if v < lo or v > hi]
    return {
        "n": n,
        "sum": float(np.sum(arr)),
        "mean": mean,
        "median": med,
        "mode": float(stats.mode(arr, keepdims=True).mode[0]) if n else float('nan'),
        "min": float(np.min(arr)),
        "max": float(np.max(arr)),
        "variance": var_pop,
        "variance_sample": var_sam,
        "std": sd_pop,
        "std_sample": sd_sam,
        "cv": float(sd_sam / mean) if mean != 0 else float('inf'),
        "range": float(np.max(arr) - np.min(arr)),
        "iqr": iqr,
        "mad": mad_v,
        "q1": q1, "q3": q3, "p10": p10, "p90": p90,
        "skewness": skew, "kurtosis": kurt,
        "geometricMean": gmean, "harmonicMean": hmean,
        "zscores_summary": {
            "mean": float(np.mean(zscores)),
            "std": float(np.std(zscores, ddof=1)) if n > 1 else float('nan')
        },
        "outliers_z_gt_3_indices": outliers_z3_idx,
        "outliers_iqr_indices": outliers_iqr_idx
    }

def correlate_block(X: NDArray, Y: NDArray, matrix_cols: Dict[str, NDArray]) -> Dict[str, Any]:
    r_pear, p_pear = stats.pearsonr(X, Y)
    r_spear, p_spear = stats.spearmanr(X, Y)
    tau, p_kend = stats.kendalltau(X, Y)

    # pairwise matrix (pearson)
    keys = list(matrix_cols.keys())
    pairs = {}
    for i in range(len(keys)):
        for j in range(i+1, len(keys)):
            a, b = matrix_cols[keys[i]], matrix_cols[keys[j]]
            r, p = stats.pearsonr(a, b)
            pairs[f"{keys[i]}|{keys[j]}"] = {"r": float(r), "p": float(p), "df": int(len(a)-2)}
    return {
        "pearson": {"r": float(r_pear), "p": float(p_pear), "df": int(len(X)-2)},
        "spearman": {"r": float(r_spear), "p": float(p_spear)},
        "kendall": {"tau": float(tau), "p": float(p_kend)},
        "matrix_pearson": pairs
    }

def compare_means_block(groups: Dict[str, NDArray], before: NDArray, after: NDArray, X: NDArray) -> Dict[str, Any]:
    # choose two groups for 2-sample
    keys = list(groups.keys())
    g0, g1 = groups[keys[0]], groups[keys[1]]

    t_eq, p_eq = stats.ttest_ind(g0, g1, equal_var=True)
    t_w, p_w = stats.ttest_ind(g0, g1, equal_var=False)
    s0 = np.var(g0, ddof=1); s1 = np.var(g1, ddof=1)
    n0 = len(g0); n1 = len(g1)
    welch_df = (s0/n0 + s1/n1)**2 / ( (s0**2)/((n0**2)*(n0-1)) + (s1**2)/((n1**2)*(n1-1)) )

    nmin = min(len(before), len(after))
    tp, pp = stats.ttest_rel(before[:nmin], after[:nmin])

    mu0 = float(np.mean(X))  # or any constant, here: mean(X) as a neutral target
    t1, p1 = stats.ttest_1samp(X, popmean=mu0)
    df1 = int(len(X) - 1)

    # ANOVA classic
    gvals = list(groups.values())
    F_classic, p_classic = stats.f_oneway(*gvals)
    k = len(gvals)
    N = sum(len(g) for g in gvals)
    dfb = k - 1
    dfw = N - k

    # Welch ANOVA
    Fw, dfw1, dfw2, pw = welch_anova(*gvals)

    return {
        "independent_student": {"t": float(t_eq), "p": float(p_eq), "df": int(n0+n1-2)},
        "independent_welch":   {"t": float(t_w),  "p": float(p_w),  "df": float(welch_df)},
        "paired":              {"t": float(tp),   "p": float(pp),   "df": int(nmin-1)},
        "one_sample":          {"t": float(t1),   "p": float(p1),   "df": df1, "mu0": mu0},
        "anova":               {"F": float(F_classic), "p": float(p_classic), "dfBetween": int(dfb), "dfWithin": int(dfw)},
        "anova_welch":         {"F": float(Fw), "p": float(pw), "dfBetween": float(dfw1), "dfWithin": float(dfw2)}
    }

def regression_block(X: NDArray, Y: NDArray, score: NDArray) -> Dict[str, Any]:
    slope, intercept, r_val, p_val, std_err = stats.linregress(X, Y)
    # simple logistic target from score>threshold (median)
    y_bin = (score > np.median(score)).astype(int)
    X1 = np.stack([X], axis=1)
    logit = try_logistic(X1, y_bin)
    return {
        "linear": {
            "slope": float(slope), "intercept": float(intercept),
            "r": float(r_val), "r2": float(r_val**2),
            "p": float(p_val), "stderr": float(std_err)
        },
        "logistic": logit
    }

def cdf_block() -> Dict[str, Any]:
    return {
        "phi0": float(stats.norm.cdf(0.0)),
        "t_2_df10": float(stats.t.cdf(2.0, df=10)),
        "f_3_df2_20": float(stats.f.cdf(3.0, dfn=2, dfd=20))
    }

def clustering_block(cols: Dict[str, NDArray]) -> Dict[str, Any]:
    names, D = pairwise_mad_distances(cols)
    # Choose moderate params; tweak as needed
    eps, minPts = 0.5, 2
    labels = dbscan_from_distance(D, eps=eps, min_pts=minPts)
    hdb = try_hdbscan(names, D, min_cluster_size=2)
    return {
        "distance_metric": "mad",
        "col_order": names,
        "distances": [listify(r) for r in D],
        "dbscan": {"eps": eps, "minPts": minPts, "labels": labels},
        "hdbscan": hdb
    }

# ---------- main compute ----------

def compute_for_dataset(ds: Dict[str, Any]) -> Dict[str, Any]:
    D = ds["data"]
    X = np.array(D["X"], dtype=float)
    Y = np.array(D["Y"], dtype=float)
    G = np.array(D["group"], dtype=int)
    S = np.array(D["score"], dtype=float)
    B = np.array(D["before"], dtype=float)
    A = np.array(D["after"], dtype=float)
    items = np.vstack([D["Q1"], D["Q2"], D["Q3"], D["Q4"]]).T
    item_names = ["Q1","Q2","Q3","Q4"]

    # Stats (descriptives) for key columns
    stats_block = {
        "X": descriptive_block(X),
        "Y": descriptive_block(Y),
        "score": descriptive_block(S),
        "before": descriptive_block(B),
        "after": descriptive_block(A)
    }

    # Correlate
    corr = correlate_block(X, Y, {"X": X, "Y": Y, "score": S})

    # CompareMeans: build 3 groups from score by G
    groups = {str(i): S[G==i] for i in range(3)}
    cmpm = compare_means_block(groups, B, A, X)

    # Regression
    regr = regression_block(X, Y, S)

    # CDFs
    cdf = cdf_block()

    # Reliability
    alpha = float(cronbach_alpha(items))
    alpha_del = cronbach_if_deleted(items, item_names)

    # Clustering over column-series
    cols_for_clu = {
        "X": X, "Y": Y, "score": S, "before": B, "after": A,
        "Q1": items[:,0], "Q2": items[:,1], "Q3": items[:,2], "Q4": items[:,3]
    }
    clu = clustering_block(cols_for_clu)

    return {
        "data":ds["data"],
        "meta": ds["meta"],
        "stats": stats_block,
        "correlate": corr,
        "compare_means": cmpm,
        "regression": regr,
        "cdf": cdf,
        "reliability": {
            "cronbach_alpha": alpha,
            "if_items_deleted": alpha_del
        },
        "clustering": clu
    }

# ---------- entry ----------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rows", type=int, default=120)
    ap.add_argument("--seed", type=int, default=12345)
    ap.add_argument("--out", type=str, default="./goldens/out")
    ap.add_argument("--sets", type=int, default=5)
    args = ap.parse_args()
    ensure_dir(args.out)
    sets = []
    for kind in range(args.sets):
        ds = make_dataset(kind, args.rows, args.seed + kind)
        res = compute_for_dataset(ds)
        fname = os.path.join(args.out, f"set-{kind:03d}.json")
        with open(fname, "w", encoding="utf-8") as f:
            json.dump(res, f, ensure_ascii=False, indent=2)
        sets.append(fname)

    print(json.dumps({"written": sets}, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    main()