"""
directrs CLI — Typer application.

Commands:
  train    Fit DirectRS model on a CSV file, generate all plots/tables.
  predict  Load a saved model and predict on new data.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Optional

import joblib
import numpy as np
import pandas as pd
import typer
from rich.console import Console
from rich.table import Table
from sklearn.metrics import (
    accuracy_score,
    brier_score_loss,
    explained_variance_score,
    f1_score,
    mean_absolute_error,
    mean_squared_error,
    r2_score,
    roc_auc_score,
)
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import label_binarize

from .models import DirectRSClassifier, DirectRSRegressor
from .plots import (
    plot_embedding_scatter,
    plot_feature_importance,
    plot_feature_loadings,
    plot_g_matrix,
    plot_metrics_comparison,
    plot_predictions_vs_actual,
    plot_stretch_spectrum,
)

app = typer.Typer(
    name="directrs",
    help="DirectRS: piecewise-linear models via decision tree geometry extraction.",
    add_completion=False,
)
console = Console()


# ---------------------------------------------------------------------------
# Metric helpers
# ---------------------------------------------------------------------------

def _classification_metrics(
    y_true: np.ndarray, y_pred: np.ndarray, y_proba: np.ndarray, classes: np.ndarray
) -> dict[str, float]:
    metrics: dict[str, float] = {}
    metrics["Accuracy"] = float(accuracy_score(y_true, y_pred))

    # AUC (macro OvR)
    try:
        if len(classes) == 2:
            metrics["AUC"] = float(roc_auc_score(y_true, y_proba[:, 1]))
        else:
            y_bin = label_binarize(y_true, classes=classes)
            metrics["AUC"] = float(
                roc_auc_score(y_bin, y_proba, multi_class="ovr", average="macro")
            )
    except Exception:
        metrics["AUC"] = float("nan")

    # Brier score (binary only; skip for multiclass)
    if len(classes) == 2:
        metrics["Brier"] = float(brier_score_loss(y_true, y_proba[:, 1]))

    # Per-class F1
    f1_vals = f1_score(y_true, y_pred, average=None, labels=classes, zero_division=0)
    for cls, f1 in zip(classes, f1_vals):
        metrics[f"F1_class{cls}"] = float(f1)
    metrics["F1_macro"] = float(f1_score(y_true, y_pred, average="macro", zero_division=0))

    return metrics


def _regression_metrics(
    y_true: np.ndarray, y_pred: np.ndarray
) -> dict[str, float]:
    return {
        "RMSE": float(np.sqrt(mean_squared_error(y_true, y_pred))),
        "MAE": float(mean_absolute_error(y_true, y_pred)),
        "R2": float(r2_score(y_true, y_pred)),
        "ExplainedVar": float(explained_variance_score(y_true, y_pred)),
    }


def _rich_metrics_table(metrics_dict: dict[str, dict[str, float]], title: str) -> None:
    table = Table(title=title, show_header=True, header_style="bold cyan")
    table.add_column("Metric", style="bold")
    for split in metrics_dict:
        table.add_column(split, justify="right")

    # Collect all metric names
    all_metrics: set[str] = set()
    for v in metrics_dict.values():
        all_metrics.update(v.keys())

    for metric in sorted(all_metrics):
        row = [metric]
        for split in metrics_dict:
            val = metrics_dict[split].get(metric, float("nan"))
            row.append(f"{val:.4f}" if not np.isnan(val) else "N/A")
        table.add_row(*row)

    console.print(table)


# ---------------------------------------------------------------------------
# JSON summary helpers
# ---------------------------------------------------------------------------

def _tier(score: float) -> str:
    """Map a normalized importance score (0–1) to a human-readable tier."""
    if score >= 0.40:
        return "dominant"
    if score >= 0.25:
        return "high"
    if score >= 0.10:
        return "moderate"
    if score >= 0.05:
        return "low"
    return "negligible"


def _interpret_feature_importances(
    feature_names: list[str],
    tree_imp: np.ndarray,
    geom_imp: np.ndarray,
) -> dict[str, str]:
    """
    Generate human-readable summaries for tree-gain and geometry importances.
    Each summary ranks features and describes their relative contribution.
    """
    def _summarise(imp: np.ndarray, label: str) -> str:
        order = np.argsort(imp)[::-1]
        parts = []
        for rank, idx in enumerate(order):
            name = feature_names[idx]
            pct = imp[idx] * 100
            tier = _tier(imp[idx])
            parts.append(f"{name} ({pct:.1f}%, {tier})")

        top_names = [feature_names[i] for i in order if imp[i] >= 0.10]
        negligible = [feature_names[i] for i in order if imp[i] < 0.05]

        summary = f"Ranked by {label}: " + ", ".join(parts) + "."
        if top_names:
            summary += f" Features with meaningful {label} contribution (≥10%): {', '.join(top_names)}."
        if negligible:
            summary += f" Features with negligible contribution (<5%): {', '.join(negligible)}."
        return summary

    return {
        "tree_gain": _summarise(tree_imp, "tree split-gain importance"),
        "geometry": _summarise(geom_imp, "geometric stretch importance"),
    }


def _interpret_interactions(
    feature_names: list[str],
    G: np.ndarray,
    top_n: int = 10,
) -> dict:
    """
    Extract and interpret pairwise feature interactions from the off-diagonal
    elements of the geometry matrix G.

    G[i,j] is large when features i and j co-occur in leaves with predictions
    far from the global mean — a direct measure of learned interaction strength.
    """
    d = len(feature_names)
    pairs = []
    for i in range(d):
        for j in range(i + 1, d):
            pairs.append((i, j, float(G[i, j])))

    if not pairs:
        return {"ranked_pairs": [], "interpretation": "No feature interactions detected (single feature or no off-diagonal signal in G)."}

    max_val = max(abs(v) for _, _, v in pairs)
    if max_val == 0:
        return {"ranked_pairs": [], "interpretation": "All off-diagonal G values are zero — no feature interactions captured by the tree geometry."}

    # Normalize and sort descending by absolute value
    pairs_norm = [(i, j, v, abs(v) / max_val) for i, j, v in pairs]
    pairs_norm.sort(key=lambda x: x[3], reverse=True)

    def _interaction_tier(norm_score: float) -> str:
        if norm_score >= 0.60:
            return "strong"
        if norm_score >= 0.30:
            return "moderate"
        if norm_score >= 0.10:
            return "weak"
        return "negligible"

    ranked = []
    for i, j, raw, norm in pairs_norm[:top_n]:
        tier = _interaction_tier(norm)
        ranked.append({
            "features": [feature_names[i], feature_names[j]],
            "raw_score": round(raw, 6),
            "normalized_score": round(norm, 4),
            "strength": tier,
        })

    # Build prose summary
    strong = [r for r in ranked if r["strength"] == "strong"]
    moderate = [r for r in ranked if r["strength"] == "moderate"]
    weak = [r for r in ranked if r["strength"] == "weak"]
    negligible_count = len(pairs_norm) - len(strong) - len(moderate) - len(weak)

    sentences = []
    if strong:
        names = "; ".join(f"{r['features'][0]} × {r['features'][1]} ({r['normalized_score']:.2f})" for r in strong)
        sentences.append(f"Strong interactions (≥0.60 normalized): {names}.")
    if moderate:
        names = "; ".join(f"{r['features'][0]} × {r['features'][1]} ({r['normalized_score']:.2f})" for r in moderate)
        sentences.append(f"Moderate interactions (0.30–0.60): {names}.")
    if weak:
        names = "; ".join(f"{r['features'][0]} × {r['features'][1]} ({r['normalized_score']:.2f})" for r in weak)
        sentences.append(f"Weak interactions (0.10–0.30): {names}.")
    if negligible_count > 0:
        sentences.append(f"{negligible_count} other feature pair(s) show negligible interaction.")
    if not sentences:
        sentences.append("No meaningful feature interactions detected.")

    sentences.append(
        "Interaction scores reflect co-occurrence of feature pairs in leaves whose predictions "
        "deviate strongly from the global mean — a higher score means those features jointly "
        "drive the model's discriminative regions."
    )

    return {
        "ranked_pairs": ranked,
        "interpretation": " ".join(sentences),
    }

def _leaf_linear_equations(model, feature_names: list[str], task: str) -> list[dict]:
    """
    For each leaf, compute the effective linear equation in original feature space.

    Inside leaf ℓ (position pos in the leaf list), the prediction is:
      Regression:       ŷ = intercept_ℓ + Σ_j coeff_ℓ[j] * x_j
      Classification:   logit_c = intercept_ℓ_c + Σ_j coeff_ℓ_c[j] * x_j

    The effective coefficient vector is recovered by projecting the embedding-space
    weights back through S':  eff_coef = S' @ w  (S' is symmetric).
    """
    S_prime = model.S_prime_
    leaf_indices = model.leaf_indices_
    d = len(feature_names)
    block_size = d + 1

    if task == "regression":
        coef = model.ridge_.coef_          # shape (L*(d+1),)
        global_intercept = float(np.atleast_1d(model.ridge_.intercept_)[0])

        leaves = []
        for pos, node_idx in enumerate(leaf_indices):
            start = pos * block_size
            intercept = global_intercept + float(coef[start])
            w = coef[start + 1 : start + block_size]          # shape (d,)
            eff_coef = S_prime @ w                             # shape (d,)
            leaves.append({
                "leaf_position": pos,
                "node_index": int(node_idx),
                "intercept": intercept,
                "coefficients": {f: float(eff_coef[j]) for j, f in enumerate(feature_names)},
            })
        return leaves

    else:  # classification
        coef = model.lr_.coef_             # shape (n_classes_or_1, L*(d+1))
        global_intercept = model.lr_.intercept_  # shape (n_classes_or_1,)
        classes = model.classes_

        leaves = []
        for pos, node_idx in enumerate(leaf_indices):
            start = pos * block_size
            per_class = {}
            for c_idx, cls in enumerate(classes):
                intercept = float(global_intercept[c_idx]) + float(coef[c_idx, start])
                w = coef[c_idx, start + 1 : start + block_size]
                eff_coef = S_prime @ w
                per_class[str(cls)] = {
                    "intercept": intercept,
                    "coefficients": {f: float(eff_coef[j]) for j, f in enumerate(feature_names)},
                }
            leaves.append({
                "leaf_position": pos,
                "node_index": int(node_idx),
                "per_class": per_class,
            })
        return leaves


def _global_linear_equation(
    model, feature_names: list[str], task: str, leaf_equations: list[dict]
) -> dict:
    """
    Sample-weighted average of per-leaf equations.

    a_j = Σ_ℓ w_ℓ · coeff_ℓ[j],  a_0 = Σ_ℓ w_ℓ · intercept_ℓ
    where w_ℓ = fraction of training samples reaching leaf ℓ.

    This is the expected local linear model over the training distribution —
    a valid global summary, not an exact representation of the piecewise model.
    """
    internal_tree = model.tree_.tree_
    n_total = internal_tree.n_node_samples[0]
    weights = np.array([
        internal_tree.n_node_samples[idx] / n_total
        for idx in model.leaf_indices_
    ])

    if task == "regression":
        intercept = float(sum(w * eq["intercept"] for w, eq in zip(weights, leaf_equations)))
        coefficients = {
            f: float(sum(w * eq["coefficients"][f] for w, eq in zip(weights, leaf_equations)))
            for f in feature_names
        }
        return {
            "description": (
                "Sample-weighted average of per-leaf equations. "
                "Approximation: f(x) ≈ a0 + Σ_j aj·xj  "
                "(exact only if all leaves share the same coefficients)"
            ),
            "intercept": intercept,
            "coefficients": coefficients,
        }

    else:  # classification
        classes = [str(c) for c in model.classes_]
        per_class = {}
        for cls in classes:
            intercept = float(sum(
                w * eq["per_class"][cls]["intercept"] for w, eq in zip(weights, leaf_equations)
            ))
            coefficients = {
                f: float(sum(
                    w * eq["per_class"][cls]["coefficients"][f]
                    for w, eq in zip(weights, leaf_equations)
                ))
                for f in feature_names
            }
            per_class[cls] = {"intercept": intercept, "coefficients": coefficients}
        return {
            "description": (
                "Sample-weighted average of per-leaf equations per class. "
                "Approximation: logit_c(x) ≈ a0_c + Σ_j aj_c·xj  "
                "(exact only if all leaves share the same coefficients)"
            ),
            "per_class": per_class,
        }


def _build_summary_json(
    model,
    task: str,
    target: str,
    feature_names: list[str],
    train_metrics: dict[str, float],
    test_metrics: dict[str, float],
    hyperparams: dict,
    n_train: int,
    n_test: int,
) -> dict:
    imps = model.feature_importances_
    tree_imp = imps["tree"]
    geom_imp = imps["geometry"]

    linear_type = "Ridge" if task == "regression" else "LogisticRegression"
    leaf_eqs = _leaf_linear_equations(model, feature_names, task)
    global_eq = _global_linear_equation(model, feature_names, task, leaf_eqs)

    leaf_desc = (
        "ŷ = intercept + Σ_j coeff[j]·x[j]"
        if task == "regression"
        else "logit_c = intercept_c + Σ_j coeff_c[j]·x[j]"
    )

    importance_interpretation = _interpret_feature_importances(feature_names, tree_imp, geom_imp)
    interaction_info = _interpret_interactions(feature_names, model.G_)

    return {
        "task": task,
        "target": target,
        "features": feature_names,
        "hyperparameters": hyperparams,
        "data_split": {"train_samples": n_train, "test_samples": n_test},
        "metrics": {"train": train_metrics, "test": test_metrics},
        "feature_importances": {
            "tree_gain": [
                {"feature": f, "importance": float(tree_imp[j])}
                for j, f in enumerate(feature_names)
            ],
            "geometry": [
                {"feature": f, "importance": float(geom_imp[j])}
                for j, f in enumerate(feature_names)
            ],
            "interpretation": importance_interpretation,
        },
        "geometry": {
            "description": "G is the d×d geometry matrix; S_prime = G^(1/2) is the stretch matrix. "
                           "The diagonal of S_prime gives the per-feature stretch (basis for geometry importances).",
            "S_prime_diagonal": {f: float(model.S_prime_[j, j]) for j, f in enumerate(feature_names)},
            "G_matrix": model.G_.tolist(),
            "S_prime_matrix": model.S_prime_.tolist(),
            "feature_interactions": interaction_info,
        },
        "linear_model": {
            "type": linear_type,
            "n_leaves": len(leaf_eqs),
            "global_equation": global_eq,
            "description": f"Per-leaf local linear equation in original feature space: {leaf_desc}",
            "leaves": leaf_eqs,
        },
    }


# ---------------------------------------------------------------------------
# train command
# ---------------------------------------------------------------------------

@app.command()
def train(
    task: str = typer.Option(..., "--task", "-t", help="'classification' or 'regression'"),
    input: Path = typer.Option(..., "--input", "-i", help="Input CSV file"),
    target: str = typer.Option(..., "--target", help="Target column name"),
    output_model: Path = typer.Option(Path("model.pkl"), "--output-model", help="Path to save fitted model"),
    max_depth: int = typer.Option(4, "--max-depth", help="Max depth of base decision tree"),
    min_samples_split: int = typer.Option(2, "--min-samples-split"),
    min_samples_leaf: int = typer.Option(1, "--min-samples-leaf"),
    C: float = typer.Option(1.0, "--C", help="LogisticRegression C (classification only)"),
    alpha: float = typer.Option(1.0, "--alpha", help="Ridge alpha (regression only)"),
    test_size: float = typer.Option(0.3, "--test-size", help="Fraction of data for test set"),
    random_state: int = typer.Option(42, "--random-state"),
    output_dir: Path = typer.Option(Path("directrs_output"), "--output-dir", help="Directory for plots and tables"),
):
    """Train a DirectRS model and generate all plots and tables."""

    # --- Validate task ---
    if task not in ("classification", "regression"):
        console.print(f"[red]Error:[/red] --task must be 'classification' or 'regression', got '{task}'")
        raise typer.Exit(1)

    # --- Load data ---
    console.print(f"[bold]Loading data from[/bold] {input}")
    df = pd.read_csv(input)
    if target not in df.columns:
        console.print(f"[red]Error:[/red] column '{target}' not found. Available: {list(df.columns)}")
        raise typer.Exit(1)

    X = df.drop(columns=[target]).values.astype(float)
    y_raw = df[target].values
    feature_names = [c for c in df.columns if c != target]

    if task == "classification":
        y = y_raw
    else:
        y = y_raw.astype(float)

    # --- Split ---
    stratify = y if task == "classification" else None
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=stratify
    )
    console.print(f"Train: {len(X_train)} samples | Test: {len(X_test)} samples")

    # --- Fit model ---
    console.print(f"[bold]Fitting DirectRS{task.capitalize()}...[/bold]")
    if task == "classification":
        model = DirectRSClassifier(
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            C=C,
            random_state=random_state,
        )
    else:
        model = DirectRSRegressor(
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            alpha=alpha,
            random_state=random_state,
        )
    model.fit(X_train, y_train)
    console.print("[green]Model fitted.[/green]")

    # --- Predictions ---
    y_train_pred = model.predict(X_train)
    y_test_pred = model.predict(X_test)

    if task == "classification":
        y_train_proba = model.predict_proba(X_train)
        y_test_proba = model.predict_proba(X_test)
        classes = model.classes_
        train_metrics = _classification_metrics(y_train, y_train_pred, y_train_proba, classes)
        test_metrics = _classification_metrics(y_test, y_test_pred, y_test_proba, classes)
    else:
        train_metrics = _regression_metrics(y_train, y_train_pred)
        test_metrics = _regression_metrics(y_test, y_test_pred)

    metrics_all = {"Train": train_metrics, "Test": test_metrics}
    _rich_metrics_table(metrics_all, title="DirectRS Metrics")

    # --- Build JSON summary ---
    _hyperparams: dict = {
        "max_depth": max_depth,
        "min_samples_split": min_samples_split,
        "min_samples_leaf": min_samples_leaf,
        "test_size": test_size,
        "random_state": random_state,
    }
    if task == "classification":
        _hyperparams["C"] = C
    else:
        _hyperparams["alpha"] = alpha

    summary = _build_summary_json(
        model=model,
        task=task,
        target=target,
        feature_names=feature_names,
        train_metrics=train_metrics,
        test_metrics=test_metrics,
        hyperparams=_hyperparams,
        n_train=len(X_train),
        n_test=len(X_test),
    )

    # --- Save JSON summary ---
    output_dir.mkdir(parents=True, exist_ok=True)
    summary_path = output_dir / "summary.json"
    with open(summary_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    console.print(f"[bold]Summary JSON saved to[/bold] {summary_path}")

    # --- Save model ---
    output_model.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump({"model": model, "feature_names": feature_names, "task": task}, output_model)
    console.print(f"[bold]Model saved to[/bold] {output_model}")

    # --- Generate plots ---
    console.print("[bold]Generating plots...[/bold]")

    figs = {
        "g_matrix": plot_g_matrix(model.G_, feature_names),
        "stretch_spectrum": plot_stretch_spectrum(model.S_prime_, feature_names),
        "feature_loadings": plot_feature_loadings(model.S_prime_, feature_names),
        "feature_importance": plot_feature_importance(model.tree_, model.S_prime_, feature_names, task),
        "embedding_scatter": plot_embedding_scatter(X_train, model.S_prime_, y_train, feature_names, task),
        "train_predictions": plot_predictions_vs_actual(y_train, y_train_pred, "Train", task),
        "test_predictions": plot_predictions_vs_actual(y_test, y_test_pred, "Test", task),
        "metrics_comparison": plot_metrics_comparison(metrics_all, ["Train", "Test"]),
    }

    for name, fig in figs.items():
        path = output_dir / f"{name}.png"
        fig.savefig(path, dpi=120, bbox_inches="tight")
        import matplotlib.pyplot as plt
        plt.close(fig)
        console.print(f"  Saved {path}")

    # --- Save tables ---
    console.print("[bold]Saving tables...[/bold]")

    pd.DataFrame({"Metric": list(train_metrics.keys()), "Value": list(train_metrics.values())}).to_csv(
        output_dir / "train_metrics.csv", index=False
    )
    pd.DataFrame({"Metric": list(test_metrics.keys()), "Value": list(test_metrics.values())}).to_csv(
        output_dir / "test_metrics.csv", index=False
    )

    # Predictions tables
    _save_predictions_csv(X_train, y_train, y_train_pred, feature_names, output_dir / "predictions_train.csv", task)
    _save_predictions_csv(X_test, y_test, y_test_pred, feature_names, output_dir / "predictions_test.csv", task)

    console.print(f"\n[bold green]Done![/bold green] All outputs saved to {output_dir}/")


def _save_predictions_csv(
    X: np.ndarray, y_true: np.ndarray, y_pred: np.ndarray,
    feature_names: list[str], path: Path, task: str
) -> None:
    df = pd.DataFrame(X, columns=feature_names)
    df["y_true"] = y_true
    df["y_pred"] = y_pred
    if task == "regression":
        df["residual"] = y_true.astype(float) - y_pred.astype(float)
    df.to_csv(path, index=False)


# ---------------------------------------------------------------------------
# predict command
# ---------------------------------------------------------------------------

@app.command()
def predict(
    model: Path = typer.Option(..., "--model", "-m", help="Path to saved model .pkl"),
    input: Path = typer.Option(..., "--input", "-i", help="Input CSV file"),
    output: Path = typer.Option(..., "--output", "-o", help="Output CSV file with predictions"),
):
    """Load a saved model and produce predictions on new data."""
    console.print(f"[bold]Loading model from[/bold] {model}")
    bundle = joblib.load(model)
    fitted_model = bundle["model"]
    feature_names = bundle["feature_names"]
    task = bundle["task"]

    console.print(f"[bold]Loading data from[/bold] {input}")
    df = pd.read_csv(input)

    # Tolerate the target column being present (just ignore it)
    X_cols = [c for c in feature_names if c in df.columns]
    X = df[X_cols].values.astype(float)

    preds = fitted_model.predict(X)

    out_df = df.copy()
    out_df["y_pred"] = preds
    if task == "classification" and hasattr(fitted_model, "predict_proba"):
        proba = fitted_model.predict_proba(X)
        for i, cls in enumerate(fitted_model.classes_):
            out_df[f"proba_{cls}"] = proba[:, i]

    output.parent.mkdir(parents=True, exist_ok=True)
    out_df.to_csv(output, index=False)
    console.print(f"[bold green]Predictions saved to[/bold green] {output}")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    app()


if __name__ == "__main__":
    main()
