from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler

try:
    from .kernel import KernelFactory
except ImportError:
    from kernel import KernelFactory


def _read_targets_dataframe(target_path: str) -> pd.DataFrame:
    try:
        return pd.read_csv(target_path)
    except Exception:
        return pd.read_csv(target_path, sep=",", engine="python")


def _resolve_target_column(targets_df: pd.DataFrame) -> str:
    if len(targets_df.columns) < 2:
        raise ValueError(
            "Target file must contain at least two columns: molecules and target values."
        )
    return str(targets_df.columns[1])


def evaluate(
    embeddings_path: str,
    target_path: str,
    splits_dir: str,
    n_splits: int = 5,
    compute_rmse: bool = True,
    kernel_name: str = "gaussian",
    results_dir: str | None = None,
):
    """Evaluate Gaussian Process regression across pre-saved splits."""
    X = np.load(embeddings_path, allow_pickle=True)

    targets_df = _read_targets_dataframe(target_path)
    target_column = _resolve_target_column(targets_df)
    y = targets_df[target_column].to_numpy()
    print(f"Loaded {len(y)} samples from target column: {target_column}.")

    kernel = KernelFactory.create(kernel_name)
    model = GaussianProcessRegressor(
        kernel=kernel,
        n_restarts_optimizer=10,
        normalize_y=True,
        alpha=1e-6,
    )

    results = []
    train_r2, test_r2 = [], []
    if compute_rmse:
        train_rmse, test_rmse = [], []

    train_dir = Path(splits_dir) / "train"
    test_dir = Path(splits_dir) / "test"

    for i in range(1, n_splits + 1):
        train_idx = np.load(train_dir / f"train_split_{i}.npy")
        test_idx = np.load(test_dir / f"test_split_{i}.npy")

        X_train, X_test = X[train_idx], X[test_idx]
        y_train, y_test = y[train_idx], y[test_idx]

        scaler = StandardScaler()
        X_train = scaler.fit_transform(X_train)
        X_test = scaler.transform(X_test)

        print(f"\nSplit {i}: Train={len(train_idx)}, Test={len(test_idx)}")

        model.fit(X_train, y_train)
        params = model.kernel_.get_params()

        y_pred_train = model.predict(X_train)
        y_pred_test = model.predict(X_test)

        r2_train = r2_score(y_train, y_pred_train)
        r2_test = r2_score(y_test, y_pred_test)
        train_r2.append(r2_train)
        test_r2.append(r2_test)

        row = {
            "split": i,
            "train_size": len(train_idx),
            "test_size": len(test_idx),
            "r2_train": r2_train,
            "r2_test": r2_test,
        }

        for key, value in params.items():
            row[key] = value

        if compute_rmse:
            rmse_train = np.sqrt(mean_squared_error(y_train, y_pred_train))
            rmse_test = np.sqrt(mean_squared_error(y_test, y_pred_test))
            train_rmse.append(rmse_train)
            test_rmse.append(rmse_test)
            row.update({"rmse_train": rmse_train, "rmse_test": rmse_test})

        results.append(row)

        print(f"R2 Train={r2_train:.3f}, R2 Test={r2_test:.3f}")
        if compute_rmse:
            print(f"RMSE Train={rmse_train:.3f}, RMSE Test={rmse_test:.3f}")

    df_results = pd.DataFrame(results)

    summary = {
        "split": "mean +/- std",
        "train_size": "-",
        "test_size": "-",
        "r2_train": f"{df_results['r2_train'].mean():.3f} +/- {df_results['r2_train'].std():.3f}",
        "r2_test": f"{df_results['r2_test'].mean():.3f} +/- {df_results['r2_test'].std():.3f}",
    }
    if compute_rmse:
        summary.update(
            {
                "rmse_train": f"{df_results['rmse_train'].mean():.3f} +/- {df_results['rmse_train'].std():.3f}",
                "rmse_test": f"{df_results['rmse_test'].mean():.3f} +/- {df_results['rmse_test'].std():.3f}",
            }
        )

    df_results = pd.concat([df_results, pd.DataFrame([summary])], ignore_index=True)

    embed_name = Path(embeddings_path).name
    model_id = embed_name.replace(".npy", "")
    if model_id:
        model_id = model_id[0].upper() + model_id[1:]
    else:
        model_id = "Unknown"

    if results_dir:
        output_dir = Path(results_dir)
    else:
        output_dir = Path("Results") / model_id
    output_dir.mkdir(parents=True, exist_ok=True)

    output_csv = output_dir / f"results_{kernel_name}.csv"
    df_results.to_csv(output_csv, index=False)

    print(f"\nResults saved to {output_csv.resolve()}")
    print(f"Kernel used: {kernel_name}\n")

    return df_results
