"""Evaluate the fitted PCA-ridge model on the independent E-MTAB-3279 study."""

from pathlib import Path

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


PROJECT_ROOT = Path(__file__).resolve().parents[1]
EXTERNAL_DIR = PROJECT_ROOT / "data" / "raw" / "E-MTAB-3279"
PROCESSED_DIR = PROJECT_ROOT / "data" / "processed"
RESULTS_DIR = PROJECT_ROOT / "results" / "external_validation"
FIGURES_DIR = RESULTS_DIR / "figures"

COUNTS_PATH = EXTERNAL_DIR / "Accessions.snp.htseq.txt"
METADATA_PATH = EXTERNAL_DIR / "E-MTAB-3279.sdrf.txt"
MODEL_PATH = PROJECT_ROOT / "models" / "pca_ridge_model.joblib"


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 metrics(actual: pd.DataFrame, predicted: np.ndarray, model_name: 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_name,
                "mse": np.mean(error**2),
                "mae": np.mean(np.abs(error)),
                "pearson": pearson,
                "spearman": spearman,
            }
        )
    return pd.DataFrame(rows)


def load_external_expression(training_genes: pd.Index) -> tuple[pd.DataFrame, pd.DataFrame]:
    counts = pd.read_csv(COUNTS_PATH, sep="\t", index_col=0)
    counts.index.name = "gene_id"
    metadata = pd.read_csv(METADATA_PATH, sep="\t", keep_default_na=False)
    metadata = metadata.drop_duplicates("Assay Name").set_index("Assay Name")

    assay_names = counts.columns.str.replace(
        ".snp.concordant_uniq.htseq", "", regex=False
    )
    # The processed matrix labels the fourth block as C2x/S2x, while the SDRF
    # metadata labels the same assays C4/S4.
    assay_names = assay_names.str.replace(r"([CS])2x$", r"\g<1>4", regex=True)
    if not assay_names.isin(metadata.index).all():
        missing = assay_names[~assay_names.isin(metadata.index)].tolist()
        raise ValueError(f"External count columns missing from metadata: {missing}")

    missing_genes = training_genes.difference(counts.index)
    if len(missing_genes):
        raise ValueError(f"External counts lack {len(missing_genes)} training genes.")

    library_sizes = counts.sum(axis=0)
    log2_cpm = np.log2(counts.div(library_sizes, axis="columns") * 1_000_000 + 1)
    log2_cpm = log2_cpm.loc[training_genes]

    sample_metadata = metadata.loc[assay_names].reset_index()
    sample_metadata["count_column"] = counts.columns.to_numpy()
    sample_metadata["condition"] = sample_metadata[
        "Factor Value[growth condition]"
    ].str.strip()
    sample_metadata["ecotype"] = sample_metadata["Factor Value[ecotype]"].str.strip()

    sample_expression = log2_cpm.T.reset_index(names="count_column").merge(
        sample_metadata[["count_column", "ecotype", "condition"]],
        on="count_column",
        validate="one_to_one",
    )
    ecotype_condition = sample_expression.drop(columns="count_column").groupby(
        ["ecotype", "condition"]
    ).mean()
    control = ecotype_condition.xs("control", level="condition").sort_index()
    drought = ecotype_condition.xs("mild drought", level="condition").sort_index()
    if not control.index.equals(drought.index):
        raise ValueError("External control and drought ecotypes do not align.")
    return control, drought - control


def internal_development_mean(training_genes: pd.Index) -> tuple[np.ndarray, pd.DataFrame]:
    expression = pd.read_csv(
        PROCESSED_DIR / "expression_log2_cpm.tsv.gz", sep="\t", index_col=0
    )
    groups = pd.read_csv(PROCESSED_DIR / "sample_groups.csv").set_index("sample_group_id")
    groups = groups.loc[expression.columns]
    splits = pd.read_csv(PROCESSED_DIR / "ecotype_splits.csv").set_index("ecotype")

    joined = expression.T.join(groups[["ecotype", "condition"]])
    means = joined.groupby(["ecotype", "condition"]).mean()
    internal_delta = means.xs("drought", level="condition") - means.xs(
        "control", level="condition"
    )
    development_ids = splits.index[splits["split"].isin(["train", "validation"])]
    return internal_delta.loc[development_ids, training_genes].mean(axis=0).to_numpy(), splits


def main() -> None:
    model = joblib.load(MODEL_PATH)
    training_genes = pd.Index(model.feature_names_in_)
    control, external_delta = load_external_expression(training_genes)
    development_mean, internal_splits = internal_development_mean(training_genes)

    model_prediction = model.predict(control)
    mean_prediction = np.tile(development_mean, (len(control), 1))
    zero_prediction = np.zeros_like(model_prediction)

    result = pd.concat(
        [
            metrics(external_delta, zero_prediction, "zero_change"),
            metrics(external_delta, mean_prediction, "development_mean"),
            metrics(external_delta, model_prediction, "pca_ridge"),
        ],
        ignore_index=True,
    )
    internal_name = result["ecotype"].replace({"Shahdara": "Sha"})
    result["relationship_to_model_development"] = internal_name.map(
        internal_splits["split"]
    ).fillna("external_only")

    summary = result.groupby("model")[["mse", "mae", "pearson", "spearman"]].agg(
        ["mean", "median", "std"]
    )
    summary.columns = ["_".join(column) for column in summary.columns]
    summary = summary.reset_index()

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    FIGURES_DIR.mkdir(parents=True, exist_ok=True)
    result.to_csv(RESULTS_DIR / "external_metrics_by_ecotype.csv", index=False)
    summary.to_csv(RESULTS_DIR / "external_summary.csv", index=False)

    sns.set_theme(style="whitegrid", context="notebook")
    plot_data = result.loc[result["model"].ne("zero_change")]
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
    sns.boxplot(data=result, x="model", y="mae", hue="model", legend=False,
                palette=["#9AA4A6", "#3D6FB4", "#C58A2B"], ax=axes[0])
    sns.stripplot(data=result, x="model", y="mae", color="#222222", size=4,
                  alpha=0.7, ax=axes[0])
    axes[0].set(title="External-study error", xlabel="Model", ylabel="MAE")
    axes[0].tick_params(axis="x", rotation=15)

    sns.barplot(data=plot_data, x="ecotype", y="pearson", hue="model",
                palette={"development_mean": "#3D6FB4", "pca_ridge": "#C58A2B"},
                ax=axes[1])
    axes[1].axhline(0, color="#333333", linewidth=0.8)
    axes[1].set(title="External correlation by ecotype", xlabel="Ecotype",
                ylabel="Pearson correlation")
    axes[1].tick_params(axis="x", rotation=30)
    axes[1].legend(title="Model", frameon=False)
    fig.tight_layout()
    fig.savefig(FIGURES_DIR / "external_validation_comparison.png", dpi=220)
    plt.close(fig)

    print(f"External ecotypes: {', '.join(control.index)}")
    print("Relationship to original split:")
    print(result.drop_duplicates("ecotype")[["ecotype", "relationship_to_model_development"]].to_string(index=False))
    print("\nExternal summary (means across six ecotypes):")
    print(summary[["model", "mse_mean", "mae_mean", "pearson_mean", "spearman_mean"]].to_string(index=False))


if __name__ == "__main__":
    main()