#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI4EVER Tabular Genomic Selection Module
-----------------------------------------
Unified script for training and evaluating ML models on genotype-phenotype data.

Adds MULTI-TRAIT (Zhiwu design):
- duplicate genotype rows per trait (long format)
- add trait indicator columns (kminus1 or onehot)
- single-output regression model learns f(genotype, trait_code) -> phenotype
- CV split is performed by ID to avoid leakage
- optional per-trait y scaling (recommended; default ON for multi-trait)

Outputs:
- model bundle
- summary.json
- pcs.csv
- viz_train.csv (single-trait) OR viz_train_long.csv + viz_train_wide.csv (multi-trait)
- predictions_long.csv + predictions_wide.csv (multi-trait prediction)
"""

import argparse
import json
import sys
from pathlib import Path
from datetime import datetime

import joblib
import numpy as np
import pandas as pd

from sklearn.model_selection import KFold, train_test_split
from sklearn.metrics import r2_score, mean_squared_error, accuracy_score
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA

from sklearn.linear_model import Ridge, LogisticRegression
from sklearn.ensemble import (
    RandomForestRegressor,
    RandomForestClassifier,
    GradientBoostingRegressor,
    GradientBoostingClassifier,
)
from sklearn.neural_network import MLPRegressor, MLPClassifier

# Keras / TensorFlow (only imported if used)
try:
    from tensorflow import keras
except Exception:
    keras = None


# ----------------------------------------------------------------------
# Logging helper: print JSON lines to stdout so the UI can parse them
# ----------------------------------------------------------------------
def log(event: str, **kwargs):
    rec = {"event": event}
    rec.update(kwargs)
    print(json.dumps(rec), flush=True)


def progress(pct: float, phase: str):
    log("progress", pct=float(pct), phase=phase)


# ----------------------------------------------------------------------
# Metrics
# ----------------------------------------------------------------------
def regression_metrics(y_true, y_pred):
    y_true = np.asarray(y_true).ravel()
    y_pred = np.asarray(y_pred).ravel()
    corr = float(np.corrcoef(y_true, y_pred)[0, 1]) if y_true.size > 1 else float("nan")
    r2 = float(r2_score(y_true, y_pred)) if y_true.size > 1 else float("nan")
    rmse = float(np.sqrt(mean_squared_error(y_true, y_pred))) if y_true.size > 0 else float("nan")
    return corr, r2, rmse


def classification_metrics(y_true, y_pred):
    y_true = np.asarray(y_true).ravel()
    y_pred = np.asarray(y_pred).ravel()
    return float(accuracy_score(y_true, y_pred))


# ----------------------------------------------------------------------
# Model factory
# ----------------------------------------------------------------------
def parse_layers(spec: str):
    spec = (spec or "").strip()
    if not spec:
        return None
    parts = [p.strip() for p in spec.split(",") if p.strip()]
    layers = []
    for p in parts:
        try:
            v = int(p)
            if v <= 0:
                continue
            layers.append(v)
        except:
            continue
    return layers if layers else None
    
def build_estimator(task: str, model_name: str, args):
    """Return an un-fitted estimator according to task & model type."""
    if model_name == "ridge":
        if task == "reg":
            return Ridge(alpha=float(args.alpha), random_state=42)
        else:
            C = 1.0 / max(float(args.alpha), 1e-6)
            return LogisticRegression(C=C, max_iter=200, n_jobs=-1)

    elif model_name == "rf":
        if task == "reg":
            return RandomForestRegressor(
                n_estimators=int(args.n_estimators),
                random_state=42,
                n_jobs=-1,
            )
        else:
            return RandomForestClassifier(
                n_estimators=int(args.n_estimators),
                random_state=42,
                n_jobs=-1,
            )

    elif model_name == "gbdt":
        if task == "reg":
            return GradientBoostingRegressor(random_state=42)
        else:
            return GradientBoostingClassifier(random_state=42)

    elif model_name == "mlp":
        layers = parse_layers(args.layers) or [int(args.hidden), int(args.hidden)]

        if task == "reg":
            return MLPRegressor(
                hidden_layer_sizes=tuple(layers),
                max_iter=int(args.epochs),
                random_state=42,
            )
        else:
            return MLPClassifier(
                hidden_layer_sizes=tuple(layers),
                max_iter=int(args.epochs),
                random_state=42,
            )

    elif model_name == "keras":
        if keras is None:
            raise RuntimeError("TensorFlow/Keras not available. Please install tensorflow.")

        def make_keras_model(input_dim: int, n_classes: int):
            inp = keras.Input(shape=(input_dim,))

            layers_ = parse_layers(args.layers) or [int(args.hidden), int(args.hidden)]
            act = str(args.activation).strip().lower()
            if act == "gelu":
                activation_fn = keras.activations.gelu
            else:
                activation_fn = act  # "relu", "tanh", etc.

            x = inp
            for u in layers_:
                x = keras.layers.Dense(int(u), activation=activation_fn)(x)
                if float(args.dropout) > 0:
                    x = keras.layers.Dropout(float(args.dropout))(x)

            # ---- Optimizer (shared) ----
            opt_name = str(args.optimizer).strip().lower()
            lr = float(args.lr)

            if opt_name == "adam":
                opt = keras.optimizers.Adam(learning_rate=lr)
            elif opt_name == "adamw":
                opt = keras.optimizers.AdamW(
                    learning_rate=lr,
                    weight_decay=float(args.weight_decay)
                )
            elif opt_name == "rmsprop":
                opt = keras.optimizers.RMSprop(learning_rate=lr)
            elif opt_name == "sgd":
                opt = keras.optimizers.SGD(learning_rate=lr, momentum=0.9)
            else:
                raise ValueError(f"Unknown optimizer: {args.optimizer}")

            # ---- Head + loss ----
            if task == "reg":
                out = keras.layers.Dense(1, activation="linear")(x)
                model = keras.Model(inp, out)
                model.compile(optimizer=opt, loss="mse")

            else:
                if n_classes <= 2:
                    out = keras.layers.Dense(1, activation="sigmoid")(x)
                    model = keras.Model(inp, out)
                    model.compile(optimizer=opt, loss="binary_crossentropy")
                else:
                    out = keras.layers.Dense(n_classes, activation="softmax")(x)
                    model = keras.Model(inp, out)
                    model.compile(optimizer=opt, loss="sparse_categorical_crossentropy")

            return model

        return make_keras_model

    else:
        raise ValueError(f"Unknown model: {model_name}")


# ----------------------------------------------------------------------
# Optional: GWAS weights (kept from your original; minimal changes)
# ----------------------------------------------------------------------
def compute_gwas_weights(feature_cols, args):
    if not args.gwas_file:
        return None, None

    try:
        gwas = pd.read_csv(args.gwas_file, sep=None, engine="python")
        gwas.columns = gwas.columns.str.strip()
    except Exception as e:
        log("warn", msg=f"gwas_read_failed: {e}")
        return None, None

    snp_col = args.gwas_snp_col
    if snp_col not in gwas.columns:
        candidates = [c for c in gwas.columns if c.lower() in {"snp", "marker", "rsid", "id"}]
        if candidates:
            snp_col = candidates[0]
        else:
            log("warn", msg=f"gwas_missing_snp_col: {args.gwas_snp_col}")
            return None, None

    value_col = args.gwas_value_col
    if value_col is None:
        # find a p-value-like column
        p_candidates = [c for c in gwas.columns if c.lower() in {"p", "pvalue", "p_value", "pval", "p-val"}]
        if p_candidates:
            value_col = p_candidates[0]
        else:
            # fallback: first numeric column not SNP
            num_cols = [c for c in gwas.columns if c != snp_col and pd.api.types.is_numeric_dtype(gwas[c])]
            if not num_cols:
                log("warn", msg="gwas_no_numeric_value_col_found")
                return None, None
            value_col = num_cols[0]

    w_map = {}
    try:
        vals = pd.to_numeric(gwas[value_col], errors="coerce").values
        snps = gwas[snp_col].astype(str).str.strip().values
        # weights = -log10(p)
        w = -np.log10(np.clip(vals.astype("float64"), 1e-300, 1.0))
        for s, ww in zip(snps, w):
            if np.isfinite(ww):
                w_map[s] = float(ww)
    except Exception as e:
        log("warn", msg=f"gwas_weight_map_failed: {e}")
        return None, None

    weights = np.ones(len(feature_cols), dtype="float32")
    missing = 0
    for i, m in enumerate(feature_cols):
        v = w_map.get(str(m).strip(), None)
        if v is None or not np.isfinite(v) or v <= 0:
            missing += 1
            continue
        weights[i] = float(v)

    pos = weights[weights > 0]
    if pos.size > 0:
        mean_pos = float(np.mean(pos))
        if np.isfinite(mean_pos) and mean_pos > 0:
            weights = weights / mean_pos

    info = {
        "weights_enabled": True,
        "weights_method": "pvalue",
        "weights_transform": "-log10(p)",
        "weights_snp_col_used": snp_col,
        "weights_value_col_used": value_col,
        "weights_missing_markers": int(missing),
        "weights_total_markers": int(len(feature_cols)),
    }
    log("status", msg="gwas_weights_ready", **info)
    return weights.astype("float32"), info


# ----------------------------------------------------------------------
# Multi-trait helpers (Zhiwu design)
# ----------------------------------------------------------------------
def parse_csv_list(s: str):
    if s is None:
        return []
    s = str(s).strip()
    if not s:
        return []
    return [x.strip() for x in s.split(",") if x.strip()]


def melt_to_long(wide_df: pd.DataFrame, id_col: str, geno_cols, trait_cols):
    long_df = wide_df.melt(
        id_vars=[id_col] + list(geno_cols),
        value_vars=list(trait_cols),
        var_name="Trait",
        value_name="y",
    )
    long_df["y"] = pd.to_numeric(long_df["y"], errors="coerce")
    long_df = long_df.dropna(subset=["y"])
    long_df["Trait"] = long_df["Trait"].astype(str)
    return long_df


def add_trait_indicators(long_df: pd.DataFrame, trait_cols, encoding: str):
    trait_cols = [str(t) for t in trait_cols]
    traits = trait_cols[:]  # preserve UI order
    long_df["Trait"] = long_df["Trait"].astype(str)

    if encoding == "onehot":
        d = pd.get_dummies(long_df["Trait"], prefix="trait")
        # ensure stable column order
        wanted = [f"trait_{t}" for t in traits]
        for c in wanted:
            if c not in d.columns:
                d[c] = 0
        d = d[wanted]
        out = pd.concat([long_df.reset_index(drop=True), d.reset_index(drop=True)], axis=1)
        indicator_cols = list(d.columns)

    elif encoding == "kminus1":
        # Use k-1 columns; last trait is baseline all-zeros
        if len(traits) < 2:
            raise ValueError("kminus1 encoding requires at least 2 traits")
        use = traits[:-1]
        cols = []
        mat = np.zeros((len(long_df), len(use)), dtype="int8")
        trait_vals = long_df["Trait"].values
        for j, t in enumerate(use):
            cols.append(f"trait_{t}")
            mat[:, j] = (trait_vals == t).astype("int8")
        d = pd.DataFrame(mat, columns=cols)
        out = pd.concat([long_df.reset_index(drop=True), d.reset_index(drop=True)], axis=1)
        indicator_cols = cols
    else:
        raise ValueError(f"Unknown trait encoding: {encoding}")

    return out, indicator_cols, traits


def fit_trait_scaler(train_long: pd.DataFrame):
    stats = train_long.groupby("Trait")["y"].agg(["mean", "std"]).reset_index()
    stats["std"] = stats["std"].replace(0, np.nan)
    stats["std"] = stats["std"].fillna(1.0)
    # dict: trait -> (mean, std)
    return {row["Trait"]: (float(row["mean"]), float(row["std"])) for _, row in stats.iterrows()}


def apply_trait_scaler(long_df: pd.DataFrame, trait_stats: dict):
    y = long_df["y"].values.astype("float32")
    traits = long_df["Trait"].values
    y_scaled = np.empty_like(y, dtype="float32")
    for i, t in enumerate(traits):
        mu, sd = trait_stats.get(str(t), (0.0, 1.0))
        if sd == 0 or not np.isfinite(sd):
            sd = 1.0
        y_scaled[i] = (y[i] - mu) / sd
    long_df = long_df.copy()
    long_df["y_scaled"] = y_scaled
    return long_df


def invert_trait_scaler(traits, y_scaled, trait_stats: dict):
    traits = np.asarray(traits).astype(str)
    y_scaled = np.asarray(y_scaled).astype("float32")
    y = np.empty_like(y_scaled, dtype="float32")
    for i, t in enumerate(traits):
        mu, sd = trait_stats.get(str(t), (0.0, 1.0))
        y[i] = y_scaled[i] * sd + mu
    return y


# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="AI4EVER Tabular Genomic Selection")

    parser.add_argument("--genotype", required=True, help="Genotype table (samples x markers)")
    parser.add_argument("--phenotype", required=True, help="Phenotype table")

    parser.add_argument("--id-col", required=True, help="ID column present in both genotype and phenotype")

    # Single-trait (legacy)
    parser.add_argument("--target", required=False, default=None, help="Phenotype column to predict (single-trait)")

    # Multi-trait (Zhiwu design)
    parser.add_argument("--multi-trait", action="store_true", help="Enable multi-trait training (Zhiwu design)")
    parser.add_argument("--trait-cols", type=str, default="", help="Comma-separated trait columns in phenotype file")
    parser.add_argument(
        "--trait-encoding",
        choices=["kminus1", "onehot"],
        default="kminus1",
        help="Trait indicator encoding: kminus1 uses k-1 columns; onehot uses k columns.",
    )
    parser.add_argument(
        "--scale-y",
        action="store_true",
        help="Scale phenotype within each trait (recommended for multi-trait).",
    )

    parser.add_argument("--model", choices=["ridge", "rf", "gbdt", "mlp", "keras"], default="ridge")
    parser.add_argument("--standardize", action="store_true", help="Standardize features")
    parser.add_argument("--cv", type=int, default=5, help="Number of CV folds (<=1 for no CV)")

    parser.add_argument(
        "--validation-mode",
        choices=["instant", "hold", "corrected"],
        default="hold",
        help="Validation strategy label (UI only; behavior is controlled by --cv)",
    )

    # Hyperparameters
    parser.add_argument("--alpha", type=float, default=1.0, help="Ridge / logistic regularization strength")
    parser.add_argument("--n-estimators", type=int, default=300, help="Number of trees (RF / GBDT)")
    parser.add_argument("--epochs", type=int, default=100, help="Epochs for MLP / Keras")
    parser.add_argument("--batch-size", type=int, default=32, help="Batch size for Keras")
    parser.add_argument("--dropout", type=float, default=0.3, help="Dropout for Keras hidden layers")
    parser.add_argument("--hidden", type=int, default=128, help="Hidden units for MLP / Keras")
    parser.add_argument("--layers",type=str, default="", help="Comma-separated hidden layer units (e.g. 256,128,64). Overrides --hidden."
    )
    parser.add_argument("--activation", type=str, default="relu", help="Keras activation for hidden layers (relu, tanh, gelu)."
    )
    parser.add_argument("--optimizer", type=str, default="adam", help="Keras optimizer: adam, adamw, rmsprop, sgd")
    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate for the optimizer")
    parser.add_argument("--weight-decay", type=float, default=0.0, help="Weight decay for AdamW (if used)")

    parser.add_argument("--predict-file", type=str, default=None, help="Optional genotype file for prediction")

    parser.add_argument("--out-dir", type=str, default="./runs", help="Output directory")

    # --- GWAS weights (optional) ---
    parser.add_argument("--gwas-file", type=str, default=None, help="Optional GWAS table")
    parser.add_argument(
        "--gwas-weight-by",
        choices=["auto", "pvalue"],
        default="pvalue",
        help="Convert GWAS p-values into marker weights using -log10(p).",
    )
    parser.add_argument("--gwas-snp-col", type=str, default="SNP", help="SNP column in GWAS file")
    parser.add_argument("--gwas-value-col", type=str, default=None, help="GWAS p-value/effect column")
    parser.add_argument("--topk", type=int, default=2000, help="For RF/GBDT with GWAS: keep top-K markers")
    

    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    progress(0.05, "Starting...")

    # defaults: if multi-trait and user didn't specify --scale-y, enable it
    if args.multi_trait and not args.scale_y:
        args.scale_y = True

    # ------------------------------------------------------------------
    # Load data
    # ------------------------------------------------------------------
    log("status", msg="loading_data")
    progress(0.10, "Loading genotype + phenotype")
    geno = pd.read_csv(args.genotype, sep=None, engine="python")
    pheno = pd.read_csv(args.phenotype, sep=None, engine="python")
    geno.columns = geno.columns.str.strip()
    pheno.columns = pheno.columns.str.strip()

    if args.id_col not in geno.columns or args.id_col not in pheno.columns:
        log("error", msg=f"id_col '{args.id_col}' not found in both files")
        sys.exit(1)

    # ------------------------------------------------------------------
    # Determine trait mode
    # ------------------------------------------------------------------
    if args.multi_trait:
        trait_cols = parse_csv_list(args.trait_cols)
        if not trait_cols:
            log("error", msg="multi-trait requires --trait-cols (comma-separated)")
            sys.exit(1)
        for t in trait_cols:
            if t not in pheno.columns:
                log("error", msg=f"trait column '{t}' not found in phenotype file")
                sys.exit(1)
        target_label = "MULTI"  # used for folder naming
        task = "reg"            # multi-trait is regression (GS)
        log("status", msg="multi_trait_enabled", trait_cols=trait_cols, encoding=args.trait_encoding, scale_y=args.scale_y)
    else:
        if not args.target:
            log("error", msg="single-trait requires --target")
            sys.exit(1)
        if args.target not in pheno.columns:
            log("error", msg=f"target '{args.target}' not found in phenotype file")
            sys.exit(1)
        target_label = args.target

    # ------------------------------------------------------------------
    # Merge wide table
    # ------------------------------------------------------------------
    progress(0.18, "Merging genotype + phenotype")
    if args.multi_trait:
        wide = pd.merge(pheno[[args.id_col] + trait_cols], geno, on=args.id_col, how="inner")
        # no global dropna here — missing per-trait handled after melt
    else:
        wide = pd.merge(pheno[[args.id_col, args.target]], geno, on=args.id_col, how="inner").dropna()

    if wide.empty:
        log("error", msg="no overlapping samples after merge")
        sys.exit(1)

    log("status", msg=f"merged_data shape {wide.shape}")

    geno_cols = [c for c in wide.columns if c not in ([args.id_col] + (trait_cols if args.multi_trait else [args.target]))]

    # numeric X
    X_wide = wide[geno_cols].apply(pd.to_numeric, errors="coerce").fillna(0)
    feature_cols = X_wide.columns

    # ------------------------------------------------------------------
    # Optional: GWAS weights / topK
    # ------------------------------------------------------------------
    feature_weights, gwas_info = compute_gwas_weights(feature_cols, args)

    if feature_weights is not None and args.model in {"rf", "gbdt"}:
        k = int(args.topk) if args.topk else 0
        k = max(0, min(k, len(feature_cols)))
        if 0 < k < len(feature_cols):
            idx = np.argsort(-feature_weights)[:k]
            X_wide = X_wide.iloc[:, idx]
            feature_cols = X_wide.columns
            feature_weights = feature_weights[idx]
            geno_cols = list(feature_cols)
            log("status", msg="topk_marker_subset_applied", topk=int(k), kept_markers=int(len(feature_cols)))

    # ------------------------------------------------------------------
    # Task detection for single-trait
    # ------------------------------------------------------------------
    class_mapping = None
    pca_model = None
    pcs_df = None

    if not args.multi_trait:
        y_raw_col = wide[args.target]
        numeric_y = pd.to_numeric(y_raw_col, errors="coerce")
        non_null_ratio = numeric_y.notnull().mean()
        unique_count = len(y_raw_col.unique())
        task = "reg" if (non_null_ratio > 0.95 and unique_count > 10) else "cls"
        log("task", kind=task, unique_count=int(unique_count), non_null_ratio=float(non_null_ratio))

        if task == "cls":
            y_cat = y_raw_col.astype("category")
            class_mapping = {int(code): str(cat) for code, cat in enumerate(y_cat.cat.categories)}
            y_raw = y_cat.astype(str)
            y_vec = y_cat.cat.codes.astype("int32")
        else:
            y_raw = numeric_y
            y_vec = numeric_y.fillna(0).astype("float32")
    else:
        y_raw = None
        y_vec = None  # created after melt

    # ------------------------------------------------------------------
    # PCA for visualization (based on genotype only)
    # ------------------------------------------------------------------
    try:
        n_components = min(10, X_wide.shape[1])
        if n_components > 0:
            pca_model = PCA(n_components=n_components, random_state=42)
            pcs = pca_model.fit_transform(X_wide.values.astype("float32"))
            pc_cols = [f"PC{i+1}" for i in range(pcs.shape[1])]
            pcs_df = pd.DataFrame(pcs, columns=pc_cols)
            pcs_df.insert(0, args.id_col, wide[args.id_col].values)
            log("status", msg=f"computed_pcs n_components={pcs.shape[1]}")
    except Exception as e:
        log("warn", msg=f"pcs_failed: {e}")

    progress(0.30, "Computing PCs")

    # ------------------------------------------------------------------
    # Split + fit helpers
    # ------------------------------------------------------------------
    def fit_one_split(X_train, X_test, y_train, y_test):
        # Standardize X per split (always for mlp/keras)
        if args.standardize or args.model in {"mlp", "keras"}:
            scaler = StandardScaler()
            X_train_s = scaler.fit_transform(X_train)
            X_test_s = scaler.transform(X_test)
        else:
            scaler = None
            X_train_s = X_train
            X_test_s = X_test

        # Apply GWAS weights for ridge/mlp/keras (tree models already subset)
        if feature_weights is not None and args.model not in {"rf", "gbdt"}:
            try:
                X_train_s = X_train_s * feature_weights
                X_test_s = X_test_s * feature_weights
            except Exception as e:
                log("warn", msg=f"apply_feature_weights_failed: {e}")

        est = build_estimator(task, args.model, args)

        if args.model == "keras":
            model = est(
                input_dim=X_train_s.shape[1],
                n_classes=len(class_mapping) if (task == "cls" and class_mapping is not None) else 1,
            )
            callbacks = [keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)]
            model.fit(
                X_train_s,
                y_train,
                epochs=int(args.epochs),
                batch_size=int(args.batch_size),
                verbose=0,
                callbacks=callbacks,
                validation_split=0.2,
            )
            if task == "reg":
                y_pred = model.predict(X_test_s, verbose=0).ravel()
                corr, r2, rmse = regression_metrics(y_test, y_pred)
                return model, scaler, y_pred, corr, r2, rmse, None
            else:
                if len(class_mapping or []) <= 2:
                    probs = model.predict(X_test_s, verbose=0).ravel()
                    y_pred = (probs >= 0.5).astype("int32")
                else:
                    probs = model.predict(X_test_s, verbose=0)
                    y_pred = probs.argmax(axis=1)
                acc = classification_metrics(y_test, y_pred)
                return model, scaler, y_pred, None, None, None, acc
        else:
            est.fit(X_train_s, y_train)
            y_pred = est.predict(X_test_s)
            if task == "reg":
                corr, r2, rmse = regression_metrics(y_test, y_pred)
                return est, scaler, y_pred, corr, r2, rmse, None
            else:
                acc = classification_metrics(y_test, y_pred)
                return est, scaler, y_pred, None, None, None, acc

    # ------------------------------------------------------------------
    # Cross-validation / holdout
    # ------------------------------------------------------------------
    mean_corr = mean_r2 = mean_rmse = None
    mean_acc = None

    progress(0.35, "Validation")

    if args.multi_trait:
        # Split by ID to avoid leakage
        ids = wide[args.id_col].astype(str).values
        unique_ids = np.unique(ids)

        def make_Xy_from_wide(wide_subset: pd.DataFrame):
            # build long + indicators (+ optional y scaling later)
            subset = wide_subset[[args.id_col] + trait_cols + geno_cols].copy()
            subset[geno_cols] = subset[geno_cols].apply(pd.to_numeric, errors="coerce").fillna(0)
            long_df = melt_to_long(subset, args.id_col, geno_cols, trait_cols)
            long_df, indicator_cols, trait_order = add_trait_indicators(long_df, trait_cols, args.trait_encoding)
            Xmat = long_df[geno_cols + indicator_cols].values.astype("float32")
            return long_df, Xmat, indicator_cols, trait_order

        cv_corrs, cv_r2s, cv_rmses = [], [], []

        if args.cv and args.cv > 1:
            kf = KFold(n_splits=args.cv, shuffle=True, random_state=42)

            for fold, (tr_i, te_i) in enumerate(kf.split(unique_ids), start=1):
                train_ids = set(unique_ids[tr_i])
                test_ids = set(unique_ids[te_i])

                wide_tr = wide[wide[args.id_col].astype(str).isin(train_ids)]
                wide_te = wide[wide[args.id_col].astype(str).isin(test_ids)]

                train_long, X_tr, indicator_cols, trait_order = make_Xy_from_wide(wide_tr)
                test_long, X_te, _, _ = make_Xy_from_wide(wide_te)

                # per-trait y scaling (train only)
                if args.scale_y:
                    stats = fit_trait_scaler(train_long)
                    train_long = apply_trait_scaler(train_long, stats)
                    test_long = apply_trait_scaler(test_long, stats)
                    y_tr = train_long["y_scaled"].values.astype("float32")
                    y_te_scaled = test_long["y_scaled"].values.astype("float32")
                else:
                    stats = None
                    y_tr = train_long["y"].values.astype("float32")
                    y_te_scaled = test_long["y"].values.astype("float32")

                est, scaler, yhat_scaled, corr_s, r2_s, rmse_s, _ = fit_one_split(X_tr, X_te, y_tr, y_te_scaled)

                # Evaluate in original units if scaled
                if stats is not None:
                    y_te = test_long["y"].values.astype("float32")
                    yhat = invert_trait_scaler(test_long["Trait"].values, yhat_scaled, stats)
                else:
                    y_te = y_te_scaled
                    yhat = np.asarray(yhat_scaled).astype("float32")

                # --- per-trait metrics only (NO pooled/global metrics) ---
                traits = test_long["Trait"].astype(str).to_numpy()

                # keep fold-level predictions for later visualization
                if fold == 1:
                    pred_rows = []  # create once lazily

                fold_metrics = {}
                for t in np.unique(traits):
                    m = (traits == t)
                    if m.sum() < 3:
                        continue
                    c_t, r2_t, rmse_t = regression_metrics(y_te[m], yhat[m])
                    fold_metrics[t] = {"corr": float(c_t), "r2": float(r2_t), "rmse": float(rmse_t)}

                # log per-trait metrics
                log("metric_by_trait", fold=int(fold), metrics=fold_metrics)

                # accumulate for summary across folds
                if fold == 1:
                    trait_cv = {}  # trait -> {"corr":[], "r2":[], "rmse":[]}

                for t, d in fold_metrics.items():
                    trait_cv.setdefault(t, {"corr": [], "r2": [], "rmse": []})
                    trait_cv[t]["corr"].append(d["corr"])
                    trait_cv[t]["r2"].append(d["r2"])
                    trait_cv[t]["rmse"].append(d["rmse"])

                # save fold predictions for plotting (Observed vs Predicted)
                pred_rows.append(pd.DataFrame({
                    args.id_col: test_long[args.id_col].astype(str).to_numpy(),
                    "Trait": traits,
                    "y_true": y_te.astype("float32"),
                    "y_pred": np.asarray(yhat).astype("float32"),
                    "fold": int(fold),
                }))

                progress(0.35 + 0.45 * (fold / float(args.cv)), f"CV fold {fold}/{args.cv}")

            # --- summary per trait ---
            summary = {}
            for t, d in trait_cv.items():
                summary[t] = {
                    "cv_corr": float(np.nanmean(d["corr"])) if d["corr"] else float("nan"),
                    "cv_r2": float(np.nanmean(d["r2"])) if d["r2"] else float("nan"),
                    "cv_rmse": float(np.nanmean(d["rmse"])) if d["rmse"] else float("nan"),
                }

            log("metric_summary_by_trait", metrics=summary)

            
            
        else:
            # Holdout by ID
            tr_ids, te_ids = train_test_split(unique_ids, test_size=0.2, random_state=42)
            wide_tr = wide[wide[args.id_col].astype(str).isin(set(tr_ids))]
            wide_te = wide[wide[args.id_col].astype(str).isin(set(te_ids))]

            train_long, X_tr, indicator_cols, trait_order = make_Xy_from_wide(wide_tr)
            test_long, X_te, _, _ = make_Xy_from_wide(wide_te)

            if args.scale_y:
                stats = fit_trait_scaler(train_long)
                train_long = apply_trait_scaler(train_long, stats)
                test_long = apply_trait_scaler(test_long, stats)
                y_tr = train_long["y_scaled"].values.astype("float32")
                y_te_scaled = test_long["y_scaled"].values.astype("float32")
            else:
                stats = None
                y_tr = train_long["y"].values.astype("float32")
                y_te_scaled = test_long["y"].values.astype("float32")

            est, scaler, yhat_scaled, _, _, _, _ = fit_one_split(X_tr, X_te, y_tr, y_te_scaled)

            if stats is not None:
                y_te = test_long["y"].values.astype("float32")
                yhat = invert_trait_scaler(test_long["Trait"].values, yhat_scaled, stats)
            else:
                y_te = y_te_scaled
                yhat = np.asarray(yhat_scaled).astype("float32")

            # --- per-trait holdout metrics (NO pooled/global metric) ---
            traits = test_long["Trait"].astype(str).to_numpy()

            summary = {}
            for t in np.unique(traits):
                m = (traits == t)
                if m.sum() < 3:
                    continue
                c_t, r2_t, rmse_t = regression_metrics(y_te[m], yhat[m])
                summary[t] = {"corr": float(c_t), "r2": float(r2_t), "rmse": float(rmse_t)}

            log("metric_summary_by_trait", metrics=summary)

    else:
        # single-trait existing behavior (row-based split is OK)
        if args.cv and args.cv > 1:
            log("status", msg=f"starting_{args.cv}fold_cv", mode=args.validation_mode)
            kf = KFold(n_splits=args.cv, shuffle=True, random_state=42)

            cv_corrs, cv_r2s, cv_rmses, cv_accs = [], [], [], []
            for fold, (tr_idx, te_idx) in enumerate(kf.split(X_wide), start=1):
                X_tr, X_te = X_wide.iloc[tr_idx].values, X_wide.iloc[te_idx].values
                y_tr, y_te = y_vec.iloc[tr_idx].values, y_vec.iloc[te_idx].values

                _, _, _, corr, r2, rmse, acc = fit_one_split(X_tr, X_te, y_tr, y_te)
                if task == "reg":
                    cv_corrs.append(corr); cv_r2s.append(r2); cv_rmses.append(rmse)
                    log("metric", fold=int(fold), corr=float(corr), r2=float(r2), rmse=float(rmse))
                else:
                    cv_accs.append(acc)
                    log("metric", fold=int(fold), acc=float(acc))
                progress(0.35 + 0.45 * (fold / float(args.cv)), f"CV fold {fold}/{args.cv}")

            if task == "reg" and cv_corrs:
                mean_corr = float(np.nanmean(cv_corrs))
                mean_r2 = float(np.nanmean(cv_r2s))
                mean_rmse = float(np.nanmean(cv_rmses))
                log("metric_summary", cv_corr=mean_corr, cv_r2=mean_r2, cv_rmse=mean_rmse)
            elif task == "cls" and cv_accs:
                mean_acc = float(np.nanmean(cv_accs))
                log("metric_summary", cv_acc=mean_acc)

        else:
            X_tr, X_te, y_tr, y_te = train_test_split(
                X_wide.values,
                y_vec.values,
                test_size=0.2,
                random_state=42,
                stratify=y_vec.values if task == "cls" else None,
            )
            _, _, _, corr, r2, rmse, acc = fit_one_split(X_tr, X_te, y_tr, y_te)
            if task == "reg":
                mean_corr, mean_r2, mean_rmse = corr, r2, rmse
                log("metric_summary", corr=mean_corr, r2=mean_r2, rmse=mean_rmse)
            else:
                mean_acc = acc
                log("metric_summary", acc=mean_acc)

    # ------------------------------------------------------------------
    # Retrain on full dataset
    # ------------------------------------------------------------------
    log("status", msg=f"training_full_{args.model}")
    progress(0.85, "Training final model on full data")

    trait_stats_full = None
    indicator_cols_full = None
    trait_order_full = None

    if args.multi_trait:
        # build full long
        full_subset = wide[[args.id_col] + trait_cols + geno_cols].copy()
        full_subset[geno_cols] = full_subset[geno_cols].apply(pd.to_numeric, errors="coerce").fillna(0)
        full_long = melt_to_long(full_subset, args.id_col, geno_cols, trait_cols)
        full_long, indicator_cols_full, trait_order_full = add_trait_indicators(full_long, trait_cols, args.trait_encoding)

        if args.scale_y:
            trait_stats_full = fit_trait_scaler(full_long)
            full_long = apply_trait_scaler(full_long, trait_stats_full)
            y_full = full_long["y_scaled"].values.astype("float32")
        else:
            y_full = full_long["y"].values.astype("float32")

        X_full_raw = full_long[geno_cols + indicator_cols_full].values.astype("float32")
    else:
        X_full_raw = X_wide.values.astype("float32")
        y_full = y_vec.values

    # full-data X scaling
    if args.standardize or args.model in {"mlp", "keras"}:
        full_scaler = StandardScaler()
        X_full = full_scaler.fit_transform(X_full_raw)
    else:
        full_scaler = None
        X_full = X_full_raw

    # Apply GWAS weights (only to marker columns; NOT to trait indicator columns)
    if feature_weights is not None and args.model not in {"rf", "gbdt"}:
        try:
            if args.multi_trait:
                # only first len(feature_cols) columns correspond to genotype markers after any topK selection
                g = len(feature_cols)
                X_full[:, :g] = X_full[:, :g] * feature_weights
            else:
                X_full = X_full * feature_weights
        except Exception as e:
            log("warn", msg=f"apply_feature_weights_full_failed: {e}")

    # train final model
    if args.model == "keras":
        est_factory = build_estimator(task, args.model, args)
        final_model = est_factory(
            input_dim=X_full.shape[1],
            n_classes=len(class_mapping) if (task == "cls" and class_mapping is not None) else 1,
        )
        callbacks = [keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)]
        final_model.fit(
            X_full,
            y_full,
            epochs=int(args.epochs),
            batch_size=int(args.batch_size),
            verbose=0,
            callbacks=callbacks,
            validation_split=0.1,
        )
    else:
        final_model = build_estimator(task, args.model, args)
        final_model.fit(X_full, y_full)

    # ------------------------------------------------------------------
    # Save outputs
    # ------------------------------------------------------------------
    progress(0.95, "Saving outputs")
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M")
    model_dir = out_dir / f"{args.model}_{target_label}_{timestamp}"
    model_dir.mkdir(parents=True, exist_ok=True)

    # Save PCs
    if pcs_df is not None:
        pcs_path = model_dir / "pcs.csv"
        pcs_df.to_csv(pcs_path, index=False)
        log("artifact", path=str(pcs_path), kind="pcs")

    # Save GWAS weights used (if any)
    if feature_weights is not None:
        try:
            w_df = pd.DataFrame({"marker": list(feature_cols), "weight": feature_weights})
            w_path = model_dir / "gwas_weights.csv"
            w_df.to_csv(w_path, index=False)
            log("artifact", path=str(w_path), kind="gwas_weights")
        except Exception as e:
            log("warn", msg=f"gwas_weights_save_failed: {e}")

    # Save model bundle
    meta = {
        "task": task,
        "id_col": args.id_col,
        "target": args.target,
        "multi_trait": bool(args.multi_trait),
        "trait_cols": trait_cols if args.multi_trait else None,
        "trait_encoding": args.trait_encoding if args.multi_trait else None,
        "indicator_cols": indicator_cols_full if args.multi_trait else None,
        "trait_order": trait_order_full if args.multi_trait else None,
        "trait_stats": trait_stats_full if args.multi_trait else None,
        "feature_cols": list(feature_cols),
        "feature_weights": feature_weights.tolist() if feature_weights is not None and args.model not in {"rf", "gbdt"} else None,
        "gwas_info": gwas_info if feature_weights is not None else None,
        "class_mapping": class_mapping,
    }

    if args.model == "keras":
        model_path = model_dir / f"model_{target_label}_{task}.keras"
        final_model.save(model_path)
        log("artifact", path=str(model_path), kind="model")
        if full_scaler is not None:
            scaler_path = model_dir / "scaler.joblib"
            joblib.dump(full_scaler, scaler_path)
            log("artifact", path=str(scaler_path), kind="scaler")

        meta_path = model_dir / "meta.json"
        with open(meta_path, "w", encoding="utf-8") as f:
            json.dump(meta, f, indent=2)
        log("artifact", path=str(meta_path), kind="meta")
    else:
        bundle = {
            "model": final_model,
            "scaler": full_scaler,
            "meta": meta,
        }
        model_path = model_dir / "model.joblib"
        joblib.dump(bundle, model_path)
        log("artifact", path=str(model_path), kind="model")

    # Summary
    summary = {
        "task": task,
        "model": args.model,
        "target": target_label,
        "multi_trait": bool(args.multi_trait),
        "cv": int(args.cv),
        "metrics": {
            "corr": mean_corr,
            "r2": mean_r2,
            "rmse": mean_rmse,
            "acc": mean_acc,
        },
        "timestamp": timestamp,
    }
    summary_path = model_dir / "summary.json"
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)
    log("artifact", path=str(summary_path), kind="summary")

    # ------------------------------------------------------------------
    # Train visualization outputs
    # ------------------------------------------------------------------
    try:
        if args.model == "keras":
            y_pred_scaled = final_model.predict(X_full, verbose=0).ravel()
        else:
            y_pred_scaled = np.asarray(final_model.predict(X_full)).ravel()

        if args.multi_trait:
            # invert scaling
            if trait_stats_full is not None:
                y_pred = invert_trait_scaler(full_long["Trait"].values, y_pred_scaled, trait_stats_full)
            else:
                y_pred = y_pred_scaled

            viz_long = full_long[[args.id_col, "Trait", "y"]].copy()
            viz_long["predicted"] = y_pred
            if pcs_df is not None:
                viz_long = viz_long.merge(pcs_df, on=args.id_col, how="left")

            viz_long_path = model_dir / "viz_train_long.csv"
            viz_long.to_csv(viz_long_path, index=False)
            log("artifact", path=str(viz_long_path), kind="viz_train_long")
            
            # Observed vs predicted (from full-data fit)
            pred_vs_obs = viz_long[[args.id_col, "Trait", "y", "predicted"]].copy()
            pred_vs_obs = pred_vs_obs.rename(columns={"y": "y_true", "predicted": "y_pred"})

            pred_vs_obs_path = model_dir / "pred_vs_obs_long.csv"
            pred_vs_obs.to_csv(pred_vs_obs_path, index=False)
            log("artifact", path=str(pred_vs_obs_path), kind="pred_vs_obs_long")

            # wide
            wide_pred = viz_long.pivot_table(index=args.id_col, columns="Trait", values="predicted", aggfunc="mean").reset_index()
            wide_pred.columns = [args.id_col] + [f"pred_{c}" for c in wide_pred.columns[1:]]
            viz_wide_path = model_dir / "viz_train_wide.csv"
            wide_pred.to_csv(viz_wide_path, index=False)
            log("artifact", path=str(viz_wide_path), kind="viz_train_wide")

        else:
            viz_df = pd.DataFrame({args.id_col: wide[args.id_col].values})
            viz_df["phenotype"] = np.asarray(y_raw)
            if task == "cls":
                viz_df["phenotype_code"] = np.asarray(y_vec.values)
            viz_df["predicted"] = y_pred_scaled
            if pcs_df is not None:
                viz_df = viz_df.merge(pcs_df, on=args.id_col, how="left")
            viz_path = model_dir / "viz_train.csv"
            viz_df.to_csv(viz_path, index=False)
            log("artifact", path=str(viz_path), kind="viz_train")
    except Exception as e:
        log("warn", msg=f"viz_train_failed: {e}")

    # ------------------------------------------------------------------
    # Prediction mode
    # ------------------------------------------------------------------
    if args.predict_file:
        log("status", msg="predict_mode", file=args.predict_file)
        Gnew = pd.read_csv(args.predict_file, sep=None, engine="python")
        Gnew.columns = Gnew.columns.str.strip()

        # Ensure ID exists for nice outputs (but allow missing)
        has_id = args.id_col in Gnew.columns

        # align genotype marker columns
        for c in feature_cols:
            if c not in Gnew.columns:
                Gnew[c] = 0
        if has_id:
            Gnew_aligned = Gnew[[args.id_col] + list(feature_cols)].copy()
        else:
            Gnew_aligned = Gnew[list(feature_cols)].copy()

        Gnew_num = Gnew_aligned.drop(columns=[args.id_col], errors="ignore")
        Gnew_num = Gnew_num.apply(pd.to_numeric, errors="coerce").fillna(0).astype("float32")

        if args.multi_trait:
            # duplicate rows per trait and add indicators
            ids_new = Gnew_aligned[args.id_col].astype(str).values if has_id else np.array([str(i) for i in range(len(Gnew_num))])

            # create long feature matrix
            rows = []
            for t in trait_order_full:
                for i in range(len(Gnew_num)):
                    rows.append((ids_new[i], t))
            pred_long = pd.DataFrame(rows, columns=[args.id_col, "Trait"])

            # repeat genotype rows
            Xrep = np.repeat(Gnew_num.values, repeats=len(trait_order_full), axis=0)

            # indicator columns
            tmp = pred_long.copy()
            tmp["y"] = 0.0  # dummy
            tmp = pd.concat([tmp, pd.DataFrame(Xrep, columns=list(feature_cols))], axis=1)
            tmp, indicator_cols, _ = add_trait_indicators(tmp, trait_order_full, args.trait_encoding)

            X_new_raw = tmp[list(feature_cols) + list(indicator_cols)].values.astype("float32")

            # scale X
            if full_scaler is not None:
                X_new = full_scaler.transform(X_new_raw)
            else:
                X_new = X_new_raw

            # apply GWAS weights only to marker part
            if feature_weights is not None and args.model not in {"rf", "gbdt"}:
                g = len(feature_cols)
                X_new[:, :g] = X_new[:, :g] * feature_weights

            if args.model == "keras":
                yhat_scaled = final_model.predict(X_new, verbose=0).ravel()
            else:
                yhat_scaled = np.asarray(final_model.predict(X_new)).ravel()

            # invert scaling
            if trait_stats_full is not None:
                yhat = invert_trait_scaler(pred_long["Trait"].values, yhat_scaled, trait_stats_full)
            else:
                yhat = yhat_scaled

            pred_long["predicted"] = yhat
            pred_long_path = model_dir / "predictions_long.csv"
            pred_long.to_csv(pred_long_path, index=False)
            log("artifact", path=str(pred_long_path), kind="predictions_long")

            pred_wide = pred_long.pivot_table(index=args.id_col, columns="Trait", values="predicted", aggfunc="mean").reset_index()
            pred_wide.columns = [args.id_col] + [f"pred_{c}" for c in pred_wide.columns[1:]]
            pred_wide_path = model_dir / "predictions_wide.csv"
            pred_wide.to_csv(pred_wide_path, index=False)
            log("artifact", path=str(pred_wide_path), kind="predictions_wide")

        else:
            # single-trait prediction
            if full_scaler is not None:
                X_new = full_scaler.transform(Gnew_num.values)
            else:
                X_new = Gnew_num.values

            if feature_weights is not None and args.model not in {"rf", "gbdt"}:
                X_new = X_new * feature_weights

            if args.model == "keras":
                preds = final_model.predict(X_new, verbose=0).ravel()
            else:
                preds = np.asarray(final_model.predict(X_new)).ravel()

            pred_df = pd.DataFrame({"pred_" + str(args.target): preds})
            pred_path = model_dir / "predictions.csv"
            pred_df.to_csv(pred_path, index=False)
            log("artifact", path=str(pred_path), kind="predictions")

    log("done", msg="finished")


if __name__ == "__main__":
    main()
