#!/usr/bin/env python3
"""Generate a reproducible UMAP / t-SNE / PCA visualisation script with parameters sized to the data.

Reads the CSV header and row count, picks feature columns (numeric, minus --drop / --label), and writes
a script that:
  1. standardises features (or L2-normalises for --metric cosine)
  2. PCA-reduces to 50 dimensions first when there are more than 50 features (speed, denoising)
  3. runs the chosen methods with sensible defaults:
       UMAP   n_neighbors = 15 (5-50: small = local detail, large = global shape), min_dist = 0.1
              (0.0 for clustering downstream), metric, random_state (or n_jobs for speed)
       t-SNE  perplexity = min(30, (n - 1) / 3), init="pca", learning_rate="auto", max_iter=1000
       PCA    2 components with explained variance in the title (the honest baseline)
  4. with --grid, a small panel of several n_neighbors / perplexity values and seeds, because single
     runs mislead
  5. colours by --label (categorical) or a numeric column, saves PNG(s) and the 2-D coordinates
     (embedding_<method>.csv) for embedding_trust.py
The generated script needs pandas, scikit-learn, matplotlib and umap-learn (for UMAP).

Usage:
    python gen_embedding_plot.py data.csv --label segment --drop id -o plot_embedding.py
    python gen_embedding_plot.py emb.csv --methods umap tsne pca --grid --metric cosine
    python gen_embedding_plot.py data.csv --methods tsne --label y --perplexity 50

Exit codes: 0 written, 2 usage or read error.
"""

import argparse
import csv
import math
import os
import sys

TEMPLATE = '''"""2-D maps of {data!r}. Generated by gen_embedding_plot.py.

Run: python {script}   (pandas, scikit-learn, matplotlib{umap_req})
Read the maps as neighbourhood pictures: cluster sizes and distances between clusters are not meaningful
in t-SNE/UMAP. Check faithfulness with embedding_trust.py on the saved coordinates.
"""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.preprocessing import StandardScaler, normalize

DATA = {data!r}
FEATURES = {features!r}
LABEL = {label!r}
METHODS = {methods!r}
METRIC = {metric!r}
SEEDS = {seeds!r}
GRID = {grid!r}

df = pd.read_csv(DATA).dropna(subset=FEATURES)
X = df[FEATURES].to_numpy(dtype=float)
X = normalize(X) if METRIC == "cosine" else StandardScaler().fit_transform(X)
if X.shape[1] > 50:                                   # denoise + speed up neighbour search
    X = PCA(n_components=50, random_state=0).fit_transform(X)
colour = df[LABEL] if LABEL else None


def run(method, param, seed):
    if method == "pca":
        p = PCA(n_components=2, random_state=seed).fit(X)
        return p.transform(X), f"PCA ({{p.explained_variance_ratio_.sum():.0%}} var)"
    if method == "tsne":
        perp = min(param, (len(X) - 1) / 3)
        Y = TSNE(n_components=2, perplexity=perp, init="pca", learning_rate="auto", max_iter=1000,
                 metric="euclidean" if METRIC != "cosine" else "cosine", random_state=seed).fit_transform(X)
        return Y, f"t-SNE perplexity={{perp:g}} seed={{seed}}"
    import umap
    Y = umap.UMAP(n_neighbors=int(param), min_dist={min_dist}, metric=METRIC, random_state=seed).fit_transform(X)
    return Y, f"UMAP n_neighbors={{int(param)}} seed={{seed}}"


def draw(ax, Y, title):
    if colour is None:
        ax.scatter(Y[:, 0], Y[:, 1], s=6, alpha=0.7)
    elif pd.api.types.is_numeric_dtype(colour) and colour.nunique() > 12:
        sc = ax.scatter(Y[:, 0], Y[:, 1], c=colour, s=6, cmap="viridis", alpha=0.8)
        plt.colorbar(sc, ax=ax, shrink=0.7)
    else:
        for lv in sorted(colour.astype(str).unique()):
            m = (colour.astype(str) == lv).to_numpy()
            ax.scatter(Y[m, 0], Y[m, 1], s=6, alpha=0.8, label=lv)
        ax.legend(markerscale=3, fontsize=7, loc="best")
    ax.set_title(title, fontsize=9)
    ax.set_xticks([]); ax.set_yticks([])


for method in METHODS:
    params = GRID.get(method, [None]) if method != "pca" else [None]
    seeds = SEEDS if method != "pca" else [0]
    fig, axes = plt.subplots(len(seeds), len(params), figsize=(4 * len(params), 4 * len(seeds)), squeeze=False)
    for r, seed in enumerate(seeds):
        for c, param in enumerate(params):
            Y, title = run(method, param, seed)
            draw(axes[r][c], Y, title)
            if r == 0 and c == len(params) // 2:
                out = pd.DataFrame(Y, columns=["dim1", "dim2"])
                if LABEL:
                    out[LABEL] = df[LABEL].to_numpy()
                out.to_csv(f"embedding_{{method}}.csv", index=False)
    fig.tight_layout()
    fig.savefig(f"embedding_{{method}}.png", dpi=130)
    print(f"saved embedding_{{method}}.png and embedding_{{method}}.csv")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a UMAP / t-SNE / PCA plotting script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv")
    ap.add_argument("--label")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--methods", nargs="+", choices=["umap", "tsne", "pca"], default=["pca", "umap", "tsne"])
    ap.add_argument("--metric", choices=["euclidean", "cosine", "manhattan"], default="euclidean")
    ap.add_argument("--perplexity", type=float)
    ap.add_argument("--n-neighbors", type=int)
    ap.add_argument("--min-dist", type=float, default=0.1)
    ap.add_argument("--grid", action="store_true", help="panels over 3 parameter values x 2 seeds")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    try:
        with open(a.csv, newline="", encoding="utf-8-sig") as fh:
            rd = csv.DictReader(fh)
            header = rd.fieldnames or []
            rows = list(rd)
    except OSError as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    if a.label and a.label not in header:
        print(f"error: label {a.label!r} not found", file=sys.stderr)
        return 2

    def numeric(c):
        try:
            [float(r[c]) for r in rows if (r[c] or "").strip()]
            return True
        except ValueError:
            return False
    feats = [c for c in header if c != a.label and c not in a.drop and numeric(c)]
    if len(feats) < 2:
        print("error: need >= 2 numeric feature columns", file=sys.stderr)
        return 2
    n = len(rows)
    perp = a.perplexity or round(min(30.0, max(5.0, math.sqrt(n) / 2)), 1)
    nn = a.n_neighbors or (15 if n > 200 else 10)
    tidy = lambda v: int(v) if float(v).is_integer() else round(v, 1)
    if a.grid:
        grid = {"tsne": [tidy(max(5, round(perp / 3))), tidy(perp), tidy(min(100, perp * 3))],
                "umap": [max(5, nn // 3), nn, min(100, nn * 3)]}
        seeds = [0, 1]
    else:
        grid = {"tsne": [tidy(perp)], "umap": [nn]}
        seeds = [0]
    text = TEMPLATE.format(data=a.csv, script=os.path.basename(a.output or "plot_embedding.py"), features=feats,
                           label=a.label, methods=a.methods, metric=a.metric, seeds=seeds, grid=grid,
                           min_dist=a.min_dist, umap_req=", umap-learn" if "umap" in a.methods else "")
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {len(feats)} features, {n} rows, methods {a.methods}, "
              f"perplexity {perp:g}, n_neighbors {nn}{', grid' if a.grid else ''}", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())