"""Orchestrates Gaussian Process regression over one or more kernels."""

import argparse
from pathlib import Path

import numpy as np
import pandas as pd

try:
    from .run_gp import run_gp_regression
    from .split_data import SplitGenerator
except ImportError:
    from run_gp import run_gp_regression
    from split_data import SplitGenerator


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 preprocess_data(embeddings_path: str, target_path: str):
    """Load and validate molecular embeddings and target properties."""
    print("\n" + "=" * 60)
    print("STEP 1: Data Preprocessing")
    print("=" * 60)

    print(f"Loading embeddings from: {embeddings_path}")
    embeddings = np.load(embeddings_path, allow_pickle=True)
    print(f"  -> Embeddings shape: {embeddings.shape}")

    print(f"Loading target properties from: {target_path}")
    targets_df = _read_targets_dataframe(target_path)
    target_column = _resolve_target_column(targets_df)

    targets = targets_df[target_column].to_numpy()
    print(f"  -> Target column: {target_column}")
    print(f"  -> Target values: {len(targets)} samples")
    print(f"  -> Target range: [{targets.min():.3f}, {targets.max():.3f}]")

    n_samples = len(targets)
    if len(embeddings) != n_samples:
        raise ValueError(f"Mismatch: {len(embeddings)} embeddings vs {n_samples} targets")

    print(f"[OK] Data preprocessing complete: {n_samples} samples validated")
    return embeddings, targets, n_samples


def prepare_splits(
    n_samples: int,
    splits_dir: str,
    n_splits: int = 5,
    train_ratio: float = 0.7,
    random_state: int = 42,
    regenerate: bool = False,
):
    """Generate or validate train/test splits."""
    print("\n" + "=" * 60)
    print("STEP 2: Train/Test Split Preparation")
    print("=" * 60)

    splits_path = Path(splits_dir)
    train_path = splits_path / "train"
    test_path = splits_path / "test"

    splits_exist = (
        train_path.exists()
        and test_path.exists()
        and len(list(train_path.glob("train_split_*.npy"))) == n_splits
        and len(list(test_path.glob("test_split_*.npy"))) == n_splits
    )

    if splits_exist and not regenerate:
        print(f"[OK] Found existing splits in: {splits_dir}")
        print(f"  -> {n_splits} splits available")
        return str(splits_path)

    print(f"Generating {n_splits} new splits...")
    print(f"  -> Train ratio: {train_ratio:.1%}")
    print(f"  -> Random state: {random_state}")

    split_generator = SplitGenerator(n_samples=n_samples, train_ratio=train_ratio)
    splits = split_generator.generate_splits(n_splits=n_splits, random_state=random_state)

    splits_path.mkdir(exist_ok=True)
    train_path.mkdir(exist_ok=True)
    test_path.mkdir(exist_ok=True)

    for i, (train_idx, test_idx) in enumerate(splits, start=1):
        np.save(train_path / f"train_split_{i}.npy", train_idx)
        np.save(test_path / f"test_split_{i}.npy", test_idx)

    print(f"[OK] Splits saved to: {splits_path.absolute()}")
    return str(splits_path)


def run_pipeline(
    embeddings_path: str = None,
    target_path: str = None,
    splits_dir: str = None,
    n_splits: int = 5,
    train_ratio: float = 0.7,
    kernel_names: list = None,
    compute_rmse: bool = True,
    regenerate_splits: bool = False,
    results_dir: str | None = None,
):
    """Execute the complete molecular property prediction pipeline."""
    project_root = Path(__file__).resolve().parent
    embeddings_path = embeddings_path or str(project_root / "Data" / "ChemAI_TED.npy")
    target_path = target_path or str(project_root / "Data" / "SELFIES_KOA.txt")
    splits_dir = splits_dir or str(project_root / "Splits")
    kernel_names = kernel_names or ["gaussian", "laplacian", "matern52"]

    print("\n" + "=" * 60)
    print("MOLECULAR PROPERTY PREDICTION PIPELINE")
    print("Gaussian Process Regression with Multiple Kernels")
    print("=" * 60)

    _, _, n_samples = preprocess_data(embeddings_path, target_path)

    splits_dir = prepare_splits(
        n_samples=n_samples,
        splits_dir=splits_dir,
        n_splits=n_splits,
        train_ratio=train_ratio,
        regenerate=regenerate_splits,
    )

    print("\n" + "=" * 60)
    print("STEP 3: Gaussian Process Regression")
    print("=" * 60)

    results = {}
    for i, kernel_name in enumerate(kernel_names, 1):
        print(f"\n[{i}/{len(kernel_names)}] Training with {kernel_name.upper()} kernel...")

        results_df = run_gp_regression(
            embeddings_path=embeddings_path,
            target_path=target_path,
            splits_dir=splits_dir,
            n_splits=n_splits,
            compute_rmse=compute_rmse,
            kernel_name=kernel_name,
            results_dir=results_dir,
        )

        results[kernel_name] = results_df

    print("\n" + "=" * 60)
    print("PIPELINE COMPLETE")
    print("=" * 60)
    print(f"[OK] Evaluated {len(kernel_names)} kernel types")
    if results_dir:
        print(f"[OK] Results saved to: {results_dir}")
    else:
        print("[OK] Results saved to: Results/")
    print("=" * 60 + "\n")

    return results


def parse_arguments():
    """Parse command-line arguments for the pipeline."""
    project_root = Path(__file__).resolve().parent

    parser = argparse.ArgumentParser(
        description="Molecular Property Prediction using Gaussian Process Regression",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--embeddings_path",
        type=str,
        default=str(project_root / "Data" / "ChemAI_TED.npy"),
        help="Path to molecular embeddings (.npy file)",
    )
    parser.add_argument(
        "--target_path",
        type=str,
        default=str(project_root / "Data" / "SELFIES_KOA.txt"),
        help="Path to target properties file (2 columns: molecules + target)",
    )
    parser.add_argument(
        "--kernel_names",
        type=str,
        nargs="+",
        default=["gaussian"],
        choices=["gaussian", "laplacian", "matern52"],
        help="Kernel types to evaluate (can specify multiple)",
    )
    parser.add_argument(
        "--compute_rmse",
        type=str,
        choices=["0", "1"],
        default="1",
        help="Compute RMSE in addition to R2 (1=yes, 0=no)",
    )
    parser.add_argument(
        "--splits_dir",
        type=str,
        default=str(project_root / "Splits"),
        help="Directory containing or for saving train/test splits",
    )
    parser.add_argument(
        "--n_splits",
        type=int,
        default=5,
        help="Number of cross-validation splits",
    )
    parser.add_argument(
        "--train_ratio",
        type=float,
        default=0.7,
        help="Proportion of data for training (0.0-1.0)",
    )
    parser.add_argument(
        "--regenerate_splits",
        action="store_true",
        help="Force regeneration of splits (ignores existing split files)",
    )
    parser.add_argument(
        "--results_dir",
        type=str,
        default=None,
        help="Directory to write GP results CSV files (default: Results/<embedding_name>).",
    )

    return parser.parse_args()


def main_single(
    embeddings_path=None,
    target_path=None,
    kernel_names=None,
    compute_rmse=True,
    splits_dir=None,
    n_splits=5,
    train_ratio=0.7,
    regenerate_splits=False,
    results_dir=None,
):
    """Entry point for CLI and GUI/programmatic calls."""
    project_root = Path(__file__).resolve().parent
    embeddings_path = embeddings_path or str(project_root / "Data" / "ChemAI_TED.npy")
    target_path = target_path or str(project_root / "Data" / "SELFIES_KOA.txt")
    splits_dir = splits_dir or str(project_root / "Splits")
    kernel_names = kernel_names or ["gaussian"]

    if isinstance(kernel_names, str):
        kernel_names = [kernel_names]

    return run_pipeline(
        embeddings_path=embeddings_path,
        target_path=target_path,
        kernel_names=kernel_names,
        compute_rmse=compute_rmse,
        splits_dir=splits_dir,
        n_splits=n_splits,
        train_ratio=train_ratio,
        regenerate_splits=regenerate_splits,
        results_dir=results_dir,
    )


if __name__ == "__main__":
    args = parse_arguments()
    main_single(
        embeddings_path=args.embeddings_path,
        target_path=args.target_path,
        kernel_names=args.kernel_names,
        compute_rmse=(args.compute_rmse == "1"),
        splits_dir=args.splits_dir,
        n_splits=args.n_splits,
        train_ratio=args.train_ratio,
        regenerate_splits=args.regenerate_splits,
        results_dir=args.results_dir,
    )
