"""Train a supervised Partial Least Squares drought-response model."""

from pathlib import Path
import json

import joblib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import pearsonr, spearmanr
from sklearn.cross_decomposition import PLSRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = PROJECT_ROOT / "data" / "processed"
RESULTS_DIR = PROJECT_ROOT / "results" / "pls"
FIGURES_DIR = RESULTS_DIR / "figures"
MODEL_DIR = PROJECT_ROOT / "models"

COMPONENT_GRID = [1, 2, 3, 5, 10]
N_INPUT_GENES = 2_000
MAX_ITERATIONS = 500


def load_data() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    expression = pd.read_csv(
        DATA_DIR / "expression_log2_cpm.tsv.gz", sep="\t", index_col=0
    )
    groups = pd.read_csv(DATA_DIR / "sample_groups.csv").set_index("sample_group_id")
    groups = groups.loc[expression.columns]
    joined = expression.T.join(groups[["ecotype", "condition"]])
    means = joined.groupby(["ecotype", "condition"]).mean()
    control = means.xs("control", level="condition").sort_index()
    drought = means.xs("drought", level="condition").sort_index()
    delta = drought - control
    splits = pd.read_csv(DATA_DIR / "ecotype_splits.csv").set_index("ecotype")
    return control, delta, splits


def select_input_genes(control: pd.DataFrame, count: int = N_INPUT_GENES) -> pd.Index:
    return control.var(axis=0).nlargest(count).index


def make_model(n_components: int) -> Pipeline:
    return Pipeline(
        [
            ("scale_inputs", StandardScaler()),
            (
                "pls",
                PLSRegression(
                    n_components=n_components,
                    scale=False,
                    max_iter=MAX_ITERATIONS,
                    tol=1e-6,
                ),
            ),
        ]
    )


def safe_correlations(actual: np.ndarray, predicted: np.ndarray) -> tuple[float, float]:
    if np.allclose(predicted, predicted[0]) or np.allclose(actual, actual[0]):
        return np.nan, np.nan
    return pearsonr(actual, predicted).statistic, spearmanr(actual, predicted).statistic


def evaluate(actual: pd.DataFrame, predicted: np.ndarray, model: str) -> pd.DataFrame:
    rows = []
    for position, ecotype in enumerate(actual.index):
        observed = actual.loc[ecotype].to_numpy()
        estimate = predicted[position]
        error = estimate - observed
        pearson, spearman = safe_correlations(observed, estimate)
        rows.append(
            {
                "ecotype": ecotype,
                "model": model,
                "mse": np.mean(error**2),
                "mae": np.mean(np.abs(error)),
                "pearson": pearson,
                "spearman": spearman,
            }
        )
    return pd.DataFrame(rows)


def summarize(metrics: pd.DataFrame) -> pd.DataFrame:
    summary = metrics.groupby("model")[["mse", "mae", "pearson", "spearman"]].agg(
        ["mean", "median", "std"]
    )
    summary.columns = ["_".join(column) for column in summary.columns]
    return summary.reset_index()


def main() -> None:
    control, delta, splits = load_data()
    train_ids = splits.index[splits["split"].eq("train")]
    validation_ids = splits.index[splits["split"].eq("validation")]
    test_ids = splits.index[splits["split"].eq("test")]

    training_genes = select_input_genes(control.loc[train_ids])
    validation_rows = []
    for components in COMPONENT_GRID:
        model = make_model(components)
        model.fit(control.loc[train_ids, training_genes], delta.loc[train_ids])
        prediction = model.predict(control.loc[validation_ids, training_genes])
        score = evaluate(delta.loc[validation_ids], prediction, "pls")
        validation_rows.append(
            {
                "n_components": components,
                "n_input_genes": len(training_genes),
                "validation_mse": score["mse"].mean(),
                "validation_mae": score["mae"].mean(),
                "validation_pearson": score["pearson"].mean(),
                "validation_spearman": score["spearman"].mean(),
                "maximum_iterations_used": int(
                    np.max(model.named_steps["pls"].n_iter_)
                ),
            }
        )

    validation = pd.DataFrame(validation_rows).sort_values(
        ["validation_mse", "validation_mae", "n_components"]
    )
    selected_components = int(validation.iloc[0]["n_components"])

    # Re-select variable inputs using all development ecotypes after the number
    # of PLS components has been fixed on validation data.
    development_ids = train_ids.append(validation_ids)
    development_genes = select_input_genes(control.loc[development_ids])
    final_model = make_model(selected_components)
    final_model.fit(
        control.loc[development_ids, development_genes],
        delta.loc[development_ids],
    )
    pls_prediction = final_model.predict(control.loc[test_ids, development_genes])

    development_mean = delta.loc[development_ids].mean(axis=0).to_numpy()
    mean_prediction = np.tile(development_mean, (len(test_ids), 1))
    zero_prediction = np.zeros_like(pls_prediction)
    pca_ridge_metrics = pd.read_csv(
        PROJECT_ROOT / "results" / "reduced_rank_ridge" / "test_metrics_by_ecotype.csv"
    ).loc[lambda frame: frame["model"].eq("pca_ridge")]

    test_metrics = pd.concat(
        [
            evaluate(delta.loc[test_ids], zero_prediction, "zero_change"),
            evaluate(delta.loc[test_ids], mean_prediction, "development_mean"),
            pca_ridge_metrics,
            evaluate(delta.loc[test_ids], pls_prediction, "pls"),
        ],
        ignore_index=True,
    )
    test_summary = summarize(test_metrics)

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    FIGURES_DIR.mkdir(parents=True, exist_ok=True)
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    validation.to_csv(RESULTS_DIR / "validation_components.csv", index=False)
    test_metrics.to_csv(RESULTS_DIR / "test_metrics_by_ecotype.csv", index=False)
    test_summary.to_csv(RESULTS_DIR / "test_summary.csv", index=False)
    pd.Series(development_genes, name="gene_id").to_csv(
        RESULTS_DIR / "selected_input_genes.csv", index=False
    )
    joblib.dump(
        {"input_genes": development_genes, "pipeline": final_model},
        MODEL_DIR / "pls_model.joblib",
    )
    (MODEL_DIR / "pls_model.json").write_text(
        json.dumps(
            {
                "model": "Partial Least Squares regression",
                "selected_components": selected_components,
                "selected_input_genes": len(development_genes),
                "input_gene_selection": "highest variance in development controls",
                "selection_metric": "mean validation MSE across ecotypes",
                "development_ecotypes": len(development_ids),
                "test_ecotypes": len(test_ids),
                "maximum_iterations_used_final_fit": int(
                    np.max(final_model.named_steps["pls"].n_iter_)
                ),
            },
            indent=2,
        ),
        encoding="utf-8",
    )

    sns.set_theme(style="whitegrid", context="notebook")
    fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.3))
    ordered_validation = validation.sort_values("n_components")
    sns.lineplot(
        data=ordered_validation,
        x="n_components",
        y="validation_mse",
        marker="o",
        color="#356859",
        ax=axes[0],
    )
    axes[0].set(
        title="PLS component selection",
        xlabel="PLS components",
        ylabel="Validation MSE",
    )

    sns.boxplot(
        data=test_metrics,
        x="model",
        y="mae",
        hue="model",
        legend=False,
        palette=["#9AA4A6", "#3D6FB4", "#C58A2B", "#7A5195"],
        ax=axes[1],
    )
    axes[1].set(title="Supervised-model test error", xlabel="Model", ylabel="MAE")
    axes[1].tick_params(axis="x", rotation=18)
    fig.tight_layout()
    fig.savefig(FIGURES_DIR / "pls_model_comparison.png", dpi=220)
    plt.close(fig)

    print("Selected using validation data:")
    print(f"Input genes: {N_INPUT_GENES}")
    print(f"PLS components: {selected_components}")
    print(f"Validation MSE: {validation.iloc[0]['validation_mse']:.6f}")
    print("\nFinal test means:")
    print(
        test_summary[["model", "mse_mean", "mae_mean", "pearson_mean", "spearman_mean"]]
        .to_string(index=False)
    )


if __name__ == "__main__":
    main()