import argparse
import contextlib
import json
import os
import pickle
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler


ENDPOINTS = {
    "late_amd": "Status_late_amd",
    "anyga": "Status_anyga",
    "nv": "Status_nv",
}

TIME_COL = "Survival_in_years"

DEFAULT_TRAIN_FOLDS = [3, 4, 5, 6, 7, 8, 9]
DEFAULT_VAL_FOLDS = [2]
DEFAULT_TEST_FOLDS = [0, 1]


def suppress_noisy_warnings():
    warnings.filterwarnings("ignore", category=UserWarning)
    warnings.filterwarnings("ignore", category=RuntimeWarning)
    warnings.filterwarnings("ignore", category=FutureWarning)

    warnings.filterwarnings("ignore", message=".*Ill-conditioned matrix.*")
    warnings.filterwarnings("ignore", message=".*Newton-Raphson failed to converge sufficiently.*")
    warnings.filterwarnings("ignore", message=".*ConvergenceWarning.*")
    warnings.filterwarnings("ignore", message=".*overflow encountered.*")
    warnings.filterwarnings("ignore", message=".*invalid value encountered.*")
    warnings.filterwarnings("ignore", message=".*matrix inversion problems.*")
    warnings.filterwarnings("ignore", message=".*delta contains nan.*")

    os.environ.setdefault("PYTHONWARNINGS", "ignore")


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--json", required=True, help="Survival JSON file.")
    parser.add_argument("--features", default=None, help="DeepSeeNet .npz feature file.")
    parser.add_argument(
        "--endpoint",
        required=True,
        choices=list(ENDPOINTS.keys()),
        help="Survival endpoint.",
    )
    parser.add_argument(
        "--feature-set",
        default="deep_clinical",
        choices=[
            "grading",
            "grading_clinical",
            "deep",
            "deep_clinical",
            "deep_clinical_genotype",
        ],
    )
    parser.add_argument("--output-dir", required=True)

    parser.add_argument("--train-folds", nargs="+", type=int, default=DEFAULT_TRAIN_FOLDS)
    parser.add_argument("--val-folds", nargs="+", type=int, default=DEFAULT_VAL_FOLDS)
    parser.add_argument("--test-folds", nargs="+", type=int, default=DEFAULT_TEST_FOLDS)

    parser.add_argument(
        "--penalizer",
        type=float,
        default=0.01,
        help="L2/elastic-net penalizer for lifelines CoxPHFitter.",
    )
    parser.add_argument(
        "--l1-ratio",
        type=float,
        default=0.0,
        help="Elastic-net L1 ratio for lifelines CoxPHFitter. 0 = pure L2.",
    )
    parser.add_argument(
        "--min-std",
        type=float,
        default=1e-6,
        help="Drop features whose train-set std is <= this value.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=None,
        help="Keep only the top-k features by univariate train-set survival ranking.",
    )
    parser.add_argument(
        "--top-k-per-block",
        type=int,
        default=None,
        help=(
            "Select top-k features within each DeepSeeNet feature block "
            "(LE_DRUS, RE_DRUS, LE_PIG, RE_PIG). "
            "Example: --top-k-per-block 4 gives 16 image features total."
        ),
    )
    parser.add_argument(
        "--show-progress",
        action="store_true",
        help="Show lifelines optimization progress.",
    )

    return parser.parse_args()


def require_lifelines():
    try:
        from lifelines import CoxPHFitter
        from lifelines.utils import concordance_index
    except ImportError as e:
        raise ImportError(
            "lifelines is required. Install with:\n\n"
            "    pip install lifelines\n"
        ) from e

    return CoxPHFitter, concordance_index


def load_json_rows(path):
    with open(path, "r") as f:
        rows = json.load(f)

    if not isinstance(rows, list):
        raise ValueError(f"Expected JSON list, got {type(rows)}")

    return pd.DataFrame(rows)


def load_deep_features(path):
    data = np.load(path, allow_pickle=True)

    features = data["features"].astype(np.float32)
    patids = data["patids"]

    feature_names = data["feature_names"] if "feature_names" in data.files else None
    if feature_names is None:
        feature_names = np.array([f"deep_{i:03d}" for i in range(features.shape[1])])

    feature_names = [str(x) for x in feature_names]

    if features.shape[1] != len(feature_names):
        raise ValueError(
            f"features has {features.shape[1]} columns but feature_names has "
            f"{len(feature_names)} entries."
        )

    print("Deep feature sanity check:")
    print(f"  shape: {features.shape}")
    print(f"  NaN:   {int(np.isnan(features).sum())}")
    print(f"  Inf:   {int(np.isinf(features).sum())}")
    print(f"  min:   {float(np.nanmin(features)):.6g}")
    print(f"  max:   {float(np.nanmax(features)):.6g}")
    print(f"  mean:  {float(np.nanmean(features)):.6g}")
    print(f"  std:   {float(np.nanstd(features)):.6g}")

    feature_df = pd.DataFrame(features, columns=feature_names)
    feature_df.insert(0, "PATID", patids)

    return feature_df


def merge_features(df, features_path):
    if features_path is None:
        raise ValueError("--features is required for deep feature sets.")

    feature_df = load_deep_features(features_path)

    before = len(df)
    df = df.merge(feature_df, on="PATID", how="inner")
    after = len(df)

    if after == 0:
        raise ValueError("No rows remained after merging JSON with feature NPZ by PATID.")

    if after < before:
        print(f"[warning] Rows dropped after feature merge: {before - after} / {before}")

    return df


def get_feature_columns(df, feature_set):
    grading_cols = ["LE_DRUS", "RE_DRUS", "LE_PIG", "RE_PIG"]
    clinical_cols = ["age", "smkever"]
    genotype_cols = ["rs1061170_CFH", "rs10490924_ARMS2", "RiskScore"]

    deep_cols = [
        c for c in df.columns
        if (
            c.startswith("LE_DRUS_")
            or c.startswith("RE_DRUS_")
            or c.startswith("LE_PIG_")
            or c.startswith("RE_PIG_")
            or c.startswith("deep_")
        )
    ]

    if feature_set == "grading":
        cols = grading_cols
    elif feature_set == "grading_clinical":
        cols = grading_cols + clinical_cols
    elif feature_set == "deep":
        cols = deep_cols
    elif feature_set == "deep_clinical":
        cols = deep_cols + clinical_cols
    elif feature_set == "deep_clinical_genotype":
        cols = deep_cols + clinical_cols + genotype_cols
    else:
        raise ValueError(f"Unknown feature set: {feature_set}")

    missing = [c for c in cols if c not in df.columns]
    if missing:
        raise ValueError(f"Missing required feature columns: {missing}")

    if len(cols) == 0:
        raise ValueError(
            f"No feature columns found for feature_set={feature_set}. "
            "Check feature_names inside the .npz."
        )

    return cols


def clean_dataframe(df, endpoint_col, feature_cols):
    required = ["PATID", "fold", TIME_COL, endpoint_col] + feature_cols
    df = df[required].copy()

    df[endpoint_col] = df[endpoint_col].astype(int)
    df[TIME_COL] = pd.to_numeric(df[TIME_COL], errors="coerce")

    for col in feature_cols:
        df[col] = pd.to_numeric(df[col], errors="coerce")

    before = len(df)
    df = df.replace([np.inf, -np.inf], np.nan)
    df = df.dropna(subset=[TIME_COL, endpoint_col] + feature_cols)
    after = len(df)

    if after < before:
        print(f"[warning] Dropped rows with missing/invalid values: {before - after} / {before}")

    return df


def split_dataframe(df, train_folds, val_folds, test_folds):
    train_df = df[df["fold"].isin(train_folds)].copy()
    val_df = df[df["fold"].isin(val_folds)].copy()
    test_df = df[df["fold"].isin(test_folds)].copy()

    if len(train_df) == 0:
        raise ValueError("Train split is empty.")
    if len(val_df) == 0:
        print("[warning] Val split is empty.")
    if len(test_df) == 0:
        print("[warning] Test split is empty.")

    return train_df, val_df, test_df


def filter_low_variance_features(train_df, feature_cols, min_std):
    std = train_df[feature_cols].std(axis=0, ddof=0)
    keep_cols = std[std > min_std].index.tolist()
    drop_cols = [c for c in feature_cols if c not in keep_cols]

    if drop_cols:
        print(
            f"[info] Dropping low-variance features: "
            f"{len(drop_cols)} / {len(feature_cols)}"
        )

    if len(keep_cols) == 0:
        raise ValueError("All features were dropped by low-variance filtering.")

    return keep_cols, drop_cols


def get_deep_feature_block(col):
    if col.startswith("LE_DRUS_"):
        return "LE_DRUS"
    if col.startswith("RE_DRUS_"):
        return "RE_DRUS"
    if col.startswith("LE_PIG_"):
        return "LE_PIG"
    if col.startswith("RE_PIG_"):
        return "RE_PIG"
    return None


def rank_features_by_univariate_cindex(train_df, endpoint_col, feature_cols):
    _, concordance_index = require_lifelines()

    times = train_df[TIME_COL].values
    events = train_df[endpoint_col].values.astype(int)

    scores = []

    for col in feature_cols:
        values = train_df[col].values.astype(float)

        if not np.all(np.isfinite(values)):
            continue

        if np.std(values) < 1e-8:
            continue

        try:
            c = concordance_index(
                event_times=times,
                predicted_scores=values,
                event_observed=events,
            )
            score = max(float(c), 1.0 - float(c))
            scores.append((col, score, float(c)))
        except Exception:
            continue

    return sorted(scores, key=lambda x: x[1], reverse=True)


def select_top_k_features(train_df, endpoint_col, feature_cols, top_k):
    if top_k is None:
        return feature_cols, []

    if top_k <= 0:
        raise ValueError("--top-k must be positive.")

    if top_k >= len(feature_cols):
        print(f"[info] --top-k {top_k} >= n_features {len(feature_cols)}; keeping all.")
        return feature_cols, []

    scores = rank_features_by_univariate_cindex(train_df, endpoint_col, feature_cols)

    if len(scores) == 0:
        raise ValueError("Could not rank any features for top-k selection.")

    selected = [x[0] for x in scores[:top_k]]
    dropped = [c for c in feature_cols if c not in selected]

    print(f"[info] Selected top {len(selected)} / {len(feature_cols)} features")
    print("[info] Top selected features:")
    for col, score, raw_c in scores[: min(10, len(scores))]:
        print(f"  {col}: score={score:.4f}, raw_c={raw_c:.4f}")

    return selected, dropped


def select_top_k_per_block_features(train_df, endpoint_col, feature_cols, top_k_per_block):
    if top_k_per_block is None:
        return feature_cols, []

    if top_k_per_block <= 0:
        raise ValueError("--top-k-per-block must be positive.")

    block_to_cols = {
        "LE_DRUS": [],
        "RE_DRUS": [],
        "LE_PIG": [],
        "RE_PIG": [],
        "OTHER": [],
    }

    for col in feature_cols:
        block = get_deep_feature_block(col)
        if block is None:
            block_to_cols["OTHER"].append(col)
        else:
            block_to_cols[block].append(col)

    selected = []

    for block in ["LE_DRUS", "RE_DRUS", "LE_PIG", "RE_PIG"]:
        scores = rank_features_by_univariate_cindex(
            train_df=train_df,
            endpoint_col=endpoint_col,
            feature_cols=block_to_cols[block],
        )

        chosen = scores[:top_k_per_block]
        chosen_cols = [x[0] for x in chosen]
        selected.extend(chosen_cols)

        print(f"[info] Selected {len(chosen_cols)} features from {block}")
        for col, score, raw_c in chosen[: min(5, len(chosen))]:
            print(f"  {col}: score={score:.4f}, raw_c={raw_c:.4f}")

    other_cols = block_to_cols["OTHER"]
    selected.extend(other_cols)

    dropped = [c for c in feature_cols if c not in selected]

    print(
        f"[info] Block-balanced selection kept {len(selected)} / {len(feature_cols)} features "
        f"including {len(other_cols)} non-deep features"
    )

    if len(selected) == 0:
        raise ValueError("Block-balanced feature selection selected no features.")

    return selected, dropped


def scale_features(train_df, val_df, test_df, feature_cols):
    scaler = StandardScaler()

    train_x = scaler.fit_transform(train_df[feature_cols].values)

    if not np.all(np.isfinite(train_x)):
        raise ValueError("Scaled train features contain NaN or Inf.")

    train_scaled = train_df.copy()
    val_scaled = val_df.copy()
    test_scaled = test_df.copy()

    train_scaled.loc[:, feature_cols] = train_x

    if len(val_scaled):
        val_x = scaler.transform(val_scaled[feature_cols].values)
        if not np.all(np.isfinite(val_x)):
            raise ValueError("Scaled val features contain NaN or Inf.")
        val_scaled.loc[:, feature_cols] = val_x

    if len(test_scaled):
        test_x = scaler.transform(test_scaled[feature_cols].values)
        if not np.all(np.isfinite(test_x)):
            raise ValueError("Scaled test features contain NaN or Inf.")
        test_scaled.loc[:, feature_cols] = test_x

    return train_scaled, val_scaled, test_scaled, scaler


def make_cox_dataframe(df, endpoint_col, feature_cols):
    cols = [TIME_COL, endpoint_col] + feature_cols
    out = df[cols].copy()
    out = out.rename(columns={TIME_COL: "duration", endpoint_col: "event"})
    return out


def fit_cox(train_df, endpoint_col, feature_cols, penalizer, l1_ratio, show_progress):
    CoxPHFitter, _ = require_lifelines()

    cox_df = make_cox_dataframe(train_df, endpoint_col, feature_cols)

    if not np.all(np.isfinite(cox_df[feature_cols].values)):
        raise ValueError("Cox training matrix contains NaN or Inf before fitting.")

    model = CoxPHFitter(
        penalizer=penalizer,
        l1_ratio=l1_ratio,
    )

    if show_progress:
        model.fit(
            cox_df,
            duration_col="duration",
            event_col="event",
            show_progress=True,
        )
    else:
        with open(os.devnull, "w") as devnull:
            with contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
                model.fit(
                    cox_df,
                    duration_col="duration",
                    event_col="event",
                    show_progress=False,
                )

    return model


def evaluate_split(model, df, endpoint_col, feature_cols, split_name):
    _, concordance_index = require_lifelines()

    if len(df) == 0:
        return {
            "split": split_name,
            "n": 0,
            "events": 0,
            "c_index": None,
        }, pd.DataFrame()

    x = df[feature_cols]
    risk_score = model.predict_partial_hazard(x).values.reshape(-1)

    if not np.all(np.isfinite(risk_score)):
        raise ValueError(f"{split_name} risk scores contain NaN or Inf.")

    c_index = concordance_index(
        event_times=df[TIME_COL].values,
        predicted_scores=-risk_score,
        event_observed=df[endpoint_col].values,
    )

    pred_df = pd.DataFrame(
        {
            "PATID": df["PATID"].values,
            "fold": df["fold"].values,
            "time": df[TIME_COL].values,
            "event": df[endpoint_col].values,
            "risk_score": risk_score,
        }
    )

    horizons = [1, 2, 3, 4, 5]
    surv = model.predict_survival_function(x, times=horizons)

    for year in horizons:
        survival_prob = surv.loc[year].values
        pred_df[f"survival_{year}y"] = survival_prob
        pred_df[f"risk_{year}y"] = 1.0 - survival_prob

    metrics = {
        "split": split_name,
        "n": int(len(df)),
        "events": int(df[endpoint_col].sum()),
        "c_index": float(c_index),
    }

    return metrics, pred_df


def save_json(obj, path):
    with open(path, "w") as f:
        json.dump(obj, f, indent=2)


def save_pickle(obj, path):
    with open(path, "wb") as f:
        pickle.dump(obj, f)


def main():
    suppress_noisy_warnings()

    args = parse_args()

    if args.top_k is not None and args.top_k_per_block is not None:
        raise ValueError("Use either --top-k or --top-k-per-block, not both.")

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

    endpoint_col = ENDPOINTS[args.endpoint]

    print(f"Endpoint: {args.endpoint} ({endpoint_col})")
    print(f"Feature set: {args.feature_set}")

    df = load_json_rows(args.json)
    print(f"Loaded rows: {len(df)}")

    if args.feature_set.startswith("deep"):
        df = merge_features(df, args.features)
        print(f"Rows after feature merge: {len(df)}")

    initial_feature_cols = get_feature_columns(df, args.feature_set)
    feature_cols = list(initial_feature_cols)
    print(f"Initial number of features: {len(feature_cols)}")

    df = clean_dataframe(df, endpoint_col, feature_cols)
    print(f"Rows after cleaning: {len(df)}")

    train_df, val_df, test_df = split_dataframe(
        df,
        train_folds=args.train_folds,
        val_folds=args.val_folds,
        test_folds=args.test_folds,
    )

    print(
        "Split sizes: "
        f"train={len(train_df)}, val={len(val_df)}, test={len(test_df)}"
    )
    print(
        "Events: "
        f"train={int(train_df[endpoint_col].sum())}, "
        f"val={int(val_df[endpoint_col].sum())}, "
        f"test={int(test_df[endpoint_col].sum())}"
    )

    feature_cols, dropped_low_variance = filter_low_variance_features(
        train_df=train_df,
        feature_cols=feature_cols,
        min_std=args.min_std,
    )
    print(f"Features after variance filter: {len(feature_cols)}")

    dropped_top_k = []
    dropped_top_k_per_block = []

    if args.top_k_per_block is not None:
        feature_cols, dropped_top_k_per_block = select_top_k_per_block_features(
            train_df=train_df,
            endpoint_col=endpoint_col,
            feature_cols=feature_cols,
            top_k_per_block=args.top_k_per_block,
        )
        print(f"Features after block-balanced top-k selection: {len(feature_cols)}")
    else:
        feature_cols, dropped_top_k = select_top_k_features(
            train_df=train_df,
            endpoint_col=endpoint_col,
            feature_cols=feature_cols,
            top_k=args.top_k,
        )
        print(f"Features after top-k selection: {len(feature_cols)}")

    train_df, val_df, test_df, scaler = scale_features(
        train_df,
        val_df,
        test_df,
        feature_cols,
    )

    model = fit_cox(
        train_df=train_df,
        endpoint_col=endpoint_col,
        feature_cols=feature_cols,
        penalizer=args.penalizer,
        l1_ratio=args.l1_ratio,
        show_progress=args.show_progress,
    )

    metrics = {}
    for split_name, split_df in [
        ("train", train_df),
        ("val", val_df),
        ("test", test_df),
    ]:
        split_metrics, split_preds = evaluate_split(
            model=model,
            df=split_df,
            endpoint_col=endpoint_col,
            feature_cols=feature_cols,
            split_name=split_name,
        )

        metrics[split_name] = split_metrics

        if len(split_preds):
            split_preds.to_csv(output_dir / f"{split_name}_predictions.csv", index=False)

    config = {
        "json": args.json,
        "features": args.features,
        "endpoint": args.endpoint,
        "endpoint_col": endpoint_col,
        "feature_set": args.feature_set,
        "n_features_initial": int(len(initial_feature_cols)),
        "n_features_final": int(len(feature_cols)),
        "feature_cols": feature_cols,
        "dropped_low_variance": dropped_low_variance,
        "dropped_top_k": dropped_top_k,
        "dropped_top_k_per_block": dropped_top_k_per_block,
        "train_folds": args.train_folds,
        "val_folds": args.val_folds,
        "test_folds": args.test_folds,
        "penalizer": args.penalizer,
        "l1_ratio": args.l1_ratio,
        "min_std": args.min_std,
        "top_k": args.top_k,
        "top_k_per_block": args.top_k_per_block,
    }

    save_pickle(model, output_dir / "cox_model.pkl")
    save_pickle(scaler, output_dir / "scaler.pkl")
    save_json(metrics, output_dir / "metrics.json")
    save_json(config, output_dir / "config.json")

    print("\nMetrics")
    print(json.dumps(metrics, indent=2))
    print(f"\nSaved outputs to: {output_dir}")


if __name__ == "__main__":
    main()