"""Create ecotype-level splits and evaluate simple drought-response baselines."""

from pathlib import Path

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]
DATA_DIR = PROJECT_ROOT / "data" / "processed"
RESULTS_DIR = PROJECT_ROOT / "results" / "baselines"
FIGURES_DIR = RESULTS_DIR / "figures"

EXPRESSION_PATH = DATA_DIR / "expression_log2_cpm.tsv.gz"
GROUPS_PATH = DATA_DIR / "sample_groups.csv"
SPLITS_PATH = DATA_DIR / "ecotype_splits.csv"

RANDOM_SEED = 42
TRAIN_FRACTION = 0.70
VALIDATION_FRACTION = 0.15

CONTROL = "#356859"
DROUGHT = "#C58A2B"
BLUE = "#3D6FB4"


def build_ecotype_matrices() -> tuple[pd.DataFrame, pd.DataFrame]:
    expression = pd.read_csv(EXPRESSION_PATH, sep="\t", index_col=0)
    groups = pd.read_csv(GROUPS_PATH).set_index("sample_group_id")
    groups = groups.loc[expression.columns]

    sample_expression = expression.T.join(groups[["ecotype", "condition"]])

    # Col-0 and Oy-0 have repeated experimental blocks. Averaging creates one
    # descriptive control vector and one drought vector for each ecotype.
    ecotype_condition = sample_expression.groupby(["ecotype", "condition"]).mean()
    control = ecotype_condition.xs("control", level="condition").sort_index()
    drought = ecotype_condition.xs("drought", level="condition").sort_index()

    if not control.index.equals(drought.index):
        raise ValueError("Control and drought ecotype sets are different.")
    if not control.columns.equals(drought.columns):
        raise ValueError("Control and drought gene sets are different.")

    delta = drought - control
    return control, delta


def make_splits(ecotypes: pd.Index) -> pd.DataFrame:
    rng = np.random.default_rng(RANDOM_SEED)
    shuffled = rng.permutation(np.asarray(sorted(ecotypes)))

    n_total = len(shuffled)
    n_train = round(TRAIN_FRACTION * n_total)
    n_validation = round(VALIDATION_FRACTION * n_total)

    split = np.empty(n_total, dtype=object)
    split[:n_train] = "train"
    split[n_train : n_train + n_validation] = "validation"
    split[n_train + n_validation :] = "test"

    splits = pd.DataFrame({"ecotype": shuffled, "split": split})
    splits["random_seed"] = RANDOM_SEED
    return splits.sort_values("ecotype").reset_index(drop=True)


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_prediction(
    actual_delta: pd.DataFrame,
    predicted_delta: np.ndarray,
    baseline: str,
    split: str,
) -> pd.DataFrame:
    records = []
    for ecotype, row in actual_delta.iterrows():
        actual = row.to_numpy()
        error = predicted_delta - actual
        pearson, spearman = safe_correlations(actual, predicted_delta)
        records.append(
            {
                "ecotype": ecotype,
                "split": split,
                "baseline": baseline,
                "mse": np.mean(error**2),
                "mae": np.mean(np.abs(error)),
                "pearson": pearson,
                "spearman": spearman,
            }
        )
    return pd.DataFrame(records)


def summarize_metrics(metrics: pd.DataFrame) -> pd.DataFrame:
    value_columns = ["mse", "mae", "pearson", "spearman"]
    summary = (
        metrics.groupby(["split", "baseline"])[value_columns]
        .agg(["mean", "median", "std"])
        .reset_index()
    )
    summary.columns = [
        "_".join(part for part in column if part)
        if isinstance(column, tuple)
        else column
        for column in summary.columns
    ]
    return summary


def plot_metric_distributions(metrics: pd.DataFrame) -> None:
    evaluation = metrics.loc[metrics["split"].isin(["validation", "test"])].copy()
    sns.set_theme(style="whitegrid", context="notebook")

    fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
    sns.boxplot(
        data=evaluation,
        x="split",
        y="mae",
        hue="baseline",
        palette={"zero_change": "#9AA4A6", "training_mean": BLUE},
        ax=axes[0],
    )
    axes[0].set(title="Absolute error by held-out ecotype", xlabel="Split", ylabel="MAE")
    axes[0].legend(title="Baseline", frameon=False)

    correlation = evaluation.loc[evaluation["baseline"].eq("training_mean")]
    correlation_long = correlation.melt(
        id_vars=["ecotype", "split"],
        value_vars=["pearson", "spearman"],
        var_name="metric",
        value_name="correlation",
    )
    sns.boxplot(
        data=correlation_long,
        x="split",
        y="correlation",
        hue="metric",
        palette={"pearson": CONTROL, "spearman": DROUGHT},
        ax=axes[1],
    )
    axes[1].axhline(0, color="#333333", linewidth=0.8)
    axes[1].set(
        title="Training-mean response correlation",
        xlabel="Split",
        ylabel="Correlation across genes",
    )
    axes[1].legend(title="Metric", frameon=False)
    fig.tight_layout()
    fig.savefig(FIGURES_DIR / "baseline_metric_distributions.png", dpi=220)
    plt.close(fig)


def plot_representative_test_example(
    delta: pd.DataFrame, metrics: pd.DataFrame, mean_response: pd.Series
) -> None:
    test_metrics = metrics.loc[
        metrics["split"].eq("test") & metrics["baseline"].eq("training_mean")
    ].sort_values("pearson")
    representative = test_metrics.iloc[len(test_metrics) // 2]["ecotype"]

    actual = delta.loc[representative]
    predicted = mean_response

    sns.set_theme(style="whitegrid", context="notebook")
    fig, ax = plt.subplots(figsize=(6.3, 5.3))
    ax.hexbin(
        actual,
        predicted,
        gridsize=55,
        mincnt=1,
        cmap="YlGnBu",
        linewidths=0,
    )
    limit = max(abs(actual).max(), abs(predicted).max())
    ax.plot([-limit, limit], [-limit, limit], linestyle="--", color="#555555")
    pearson_value = test_metrics.loc[
        test_metrics["ecotype"].eq(representative), "pearson"
    ].iloc[0]
    ax.set(
        title=f"Representative test ecotype: {representative}",
        xlabel="Actual drought - control expression",
        ylabel="Training-mean predicted response",
    )
    ax.text(
        0.03,
        0.97,
        f"Pearson r = {pearson_value:.3f}",
        transform=ax.transAxes,
        va="top",
        bbox={"facecolor": "white", "alpha": 0.85, "edgecolor": "none"},
    )
    fig.tight_layout()
    fig.savefig(FIGURES_DIR / "representative_test_prediction.png", dpi=220)
    plt.close(fig)


def main() -> None:
    _, delta = build_ecotype_matrices()
    splits = make_splits(delta.index)

    split_lookup = splits.set_index("ecotype")["split"]
    train_ecotypes = split_lookup.index[split_lookup.eq("train")]
    mean_response = delta.loc[train_ecotypes].mean(axis=0)
    zero_response = np.zeros(delta.shape[1], dtype=float)

    metric_frames = []
    for split_name in ["train", "validation", "test"]:
        ecotypes = split_lookup.index[split_lookup.eq(split_name)]
        actual = delta.loc[ecotypes]
        metric_frames.append(
            evaluate_prediction(actual, zero_response, "zero_change", split_name)
        )
        metric_frames.append(
            evaluate_prediction(
                actual, mean_response.to_numpy(), "training_mean", split_name
            )
        )

    metrics = pd.concat(metric_frames, ignore_index=True)
    summary = summarize_metrics(metrics)

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    FIGURES_DIR.mkdir(parents=True, exist_ok=True)
    splits.to_csv(SPLITS_PATH, index=False)
    metrics.to_csv(RESULTS_DIR / "baseline_metrics_by_ecotype.csv", index=False)
    summary.to_csv(RESULTS_DIR / "baseline_summary.csv", index=False)
    mean_response.rename("mean_delta").to_csv(
        RESULTS_DIR / "training_mean_response.tsv.gz",
        sep="\t",
        compression="gzip",
    )

    plot_metric_distributions(metrics)
    plot_representative_test_example(delta, metrics, mean_response)

    print("Split sizes:")
    print(splits["split"].value_counts().reindex(["train", "validation", "test"]))
    print("\nValidation and test summary (means across ecotypes):")
    columns = ["split", "baseline", "mse_mean", "mae_mean", "pearson_mean", "spearman_mean"]
    print(summary.loc[summary["split"].ne("train"), columns].to_string(index=False))
    print(f"\nWrote outputs to {RESULTS_DIR}")


if __name__ == "__main__":
    main()