"""
Evaluation and Diagnostics Suite for GNNs and Classical Baselines.
Generates parity plots, ROC/PR curves, loss trajectories, t-SNE latent embeddings,
and compiles results/REPORT.md.
"""
import os
import sys

# Ensure project root is in sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import json
import logging
from typing import Dict, List, Any, Optional, Tuple

import yaml
import numpy as np
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE
from sklearn.metrics import (
    mean_squared_error, mean_absolute_error, r2_score,
    roc_curve, precision_recall_curve, roc_auc_score, average_precision_score
)

from src.data import get_dataloaders, NODE_DIM, EDGE_DIM
from src.models.gnn import build_gnn_model
from src.models.baseline import ClassicalFingerprintBaseline, extract_features

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# Matplotlib styling for publication quality
plt.style.use('seaborn-v0_8-whitegrid' if 'seaborn-v0_8-whitegrid' in plt.style.available else 'default')
plt.rcParams['font.sans-serif'] = 'DejaVu Sans'
plt.rcParams['axes.edgecolor'] = '#cccccc'
plt.rcParams['axes.linewidth'] = 0.8


def load_trained_model(
    model_name: str,
    dataset_name: str,
    checkpoint_dir: str = 'results/checkpoints'
) -> Tuple[Any, Optional[Dict[str, Any]]]:
    """Load trained model checkpoint from disk."""
    model_name = model_name.lower()
    dataset_name = dataset_name.lower()

    if model_name == 'baseline':
        if os.environ.get("LOW_MEMORY_MODE") == "1" and dataset_name == "tox21":
            # Avoid deserializing 650MB Tox21 Random Forest pickle on 512MB cloud free-tier VMs
            raise RuntimeError("Tox21 RF baseline disabled in LOW_MEMORY_MODE to preserve RAM limit.")
        ckpt_path = os.path.join(checkpoint_dir, f"{dataset_name}_baseline_best.pkl")
        if not os.path.exists(ckpt_path):
            raise FileNotFoundError(f"Baseline checkpoint not found at {ckpt_path}")
        return ClassicalFingerprintBaseline.load(ckpt_path), None
    else:
        ckpt_path = os.path.join(checkpoint_dir, f"{dataset_name}_{model_name}_best.pt")
        if not os.path.exists(ckpt_path):
            raise FileNotFoundError(f"GNN checkpoint not found at {ckpt_path}")
        payload = torch.load(ckpt_path, weights_only=False)
        kwargs = payload['model_kwargs']
        model = build_gnn_model(model_name=model_name, **kwargs)
        model.load_state_dict(payload['model_state_dict'])
        model.eval()
        return model, payload


def plot_loss_trajectories(dataset_name: str, history_dir: str = 'results/history', save_dir: str = 'results/plots'):
    """Plot training loss curves across GNN architectures."""
    os.makedirs(save_dir, exist_ok=True)
    models = ['gcn', 'gat', 'gin', 'mpnn']
    colors = {'gcn': '#1f77b4', 'gat': '#ff7f0e', 'gin': '#2ca02c', 'mpnn': '#d62728'}

    plt.figure(figsize=(8, 5))
    has_data = False

    for m in models:
        hist_path = os.path.join(history_dir, f"{dataset_name}_{m}_history.json")
        if os.path.exists(hist_path):
            with open(hist_path, 'r') as f:
                hist = json.load(f)
            losses = hist.get('train_loss', [])
            if losses:
                plt.plot(range(1, len(losses) + 1), losses, label=m.upper(), color=colors.get(m, '#333333'), lw=2.0)
                has_data = True

    if has_data:
        plt.xlabel("Epoch", fontsize=11, fontweight='bold')
        plt.ylabel("Training Loss", fontsize=11, fontweight='bold')
        plt.title(f"Training Loss Trajectory on {dataset_name.upper()} (CPU-Optimized)", fontsize=12, fontweight='bold')
        plt.legend(frameon=True)
        plt.tight_layout()
        plot_path = os.path.join(save_dir, f"{dataset_name}_loss_curves.png")
        plt.savefig(plot_path, dpi=300)
        plt.close()
        logger.info(f"Saved loss trajectory plot -> {plot_path}")


def plot_parity(
    dataset_name: str,
    model_name: str = 'mpnn',
    checkpoint_dir: str = 'results/checkpoints',
    save_dir: str = 'results/plots'
):
    """Generate Parity Plot (Predicted vs Actual) for ESOL Regression."""
    os.makedirs(save_dir, exist_ok=True)
    _, _, test_loader, meta = get_dataloaders(dataset_name, batch_size=32)
    test_data = meta['test']

    model, _ = load_trained_model(model_name, dataset_name, checkpoint_dir)

    if model_name == 'baseline':
        preds = model.predict(test_data).flatten()
        _, y_true = extract_features(test_data)
        y_true = y_true.flatten()
    else:
        preds_list, targets_list = [], []
        with torch.no_grad():
            for batch in test_loader:
                out = model(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch)
                preds_list.append(out.view(-1).numpy())
                targets_list.append(batch.y.view(-1).numpy())
        preds = np.concatenate(preds_list)
        y_true = np.concatenate(targets_list)

    rmse = np.sqrt(mean_squared_error(y_true, preds))
    r2 = r2_score(y_true, preds)

    plt.figure(figsize=(6.5, 6))
    plt.scatter(y_true, preds, alpha=0.7, color='#2b5c8f', edgecolors='k', linewidth=0.5, s=45, label='Test Compounds')

    min_val = min(y_true.min(), preds.min()) - 0.5
    max_val = max(y_true.max(), preds.max()) + 0.5
    plt.plot([min_val, max_val], [min_val, max_val], 'r--', lw=1.8, label='Ideal 1:1 Parity')

    plt.xlim(min_val, max_val)
    plt.ylim(min_val, max_val)
    plt.xlabel("Measured Aqueous Solubility logS (mol/L)", fontsize=11, fontweight='bold')
    plt.ylabel(f"Predicted logS by {model_name.upper()} (mol/L)", fontsize=11, fontweight='bold')
    plt.title(f"ESOL Parity Plot ({model_name.upper()} on Scaffold Test Split)\nRMSE: {rmse:.3f} | R²: {r2:.3f}", fontsize=12, fontweight='bold')
    plt.legend(frameon=True)
    plt.tight_layout()

    plot_path = os.path.join(save_dir, f"{dataset_name}_parity_{model_name}.png")
    plt.savefig(plot_path, dpi=300)
    plt.close()
    logger.info(f"Saved parity plot -> {plot_path}")


def plot_roc_pr_curves(
    dataset_name: str,
    checkpoint_dir: str = 'results/checkpoints',
    save_dir: str = 'results/plots'
):
    """Generate ROC and PR Curves for BBBP or Tox21 classification."""
    os.makedirs(save_dir, exist_ok=True)
    models = ['baseline', 'gcn', 'gat', 'gin', 'mpnn']
    colors = {'baseline': '#7f7f7f', 'gcn': '#1f77b4', 'gat': '#ff7f0e', 'gin': '#2ca02c', 'mpnn': '#d62728'}

    _, _, test_loader, meta = get_dataloaders(dataset_name, batch_size=32)
    test_data = meta['test']

    fig, (ax_roc, ax_pr) = plt.subplots(1, 2, figsize=(13, 5.5))

    for m in models:
        try:
            model, _ = load_trained_model(m, dataset_name, checkpoint_dir)
        except Exception:
            continue

        if m == 'baseline':
            preds = model.predict(test_data)[:, 0]
            _, y_true = extract_features(test_data)
            y_true = y_true[:, 0]
        else:
            preds_list, targets_list = [], []
            with torch.no_grad():
                for batch in test_loader:
                    out = torch.sigmoid(model(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch))
                    preds_list.append(out[:, 0].numpy())
                    targets_list.append(batch.y[:, 0].numpy())
            preds = np.concatenate(preds_list)
            y_true = np.concatenate(targets_list)

        mask = ~np.isnan(y_true)
        y_eval = y_true[mask]
        p_eval = preds[mask]

        if len(np.unique(y_eval)) > 1:
            fpr, tpr, _ = roc_curve(y_eval, p_eval)
            roc_auc = roc_auc_score(y_eval, p_eval)
            prec, rec, _ = precision_recall_curve(y_eval, p_eval)
            pr_auc = average_precision_score(y_eval, p_eval)

            c = colors.get(m, '#333333')
            ax_roc.plot(fpr, tpr, label=f"{m.upper()} (AUC = {roc_auc:.3f})", color=c, lw=2.0)
            ax_pr.plot(rec, prec, label=f"{m.upper()} (AUC = {pr_auc:.3f})", color=c, lw=2.0)

    ax_roc.plot([0, 1], [0, 1], 'k--', lw=1.2, label='Random Chance')
    ax_roc.set_xlabel("False Positive Rate", fontsize=11, fontweight='bold')
    ax_roc.set_ylabel("True Positive Rate", fontsize=11, fontweight='bold')
    ax_roc.set_title(f"ROC Curves on {dataset_name.upper()} Test Split", fontsize=12, fontweight='bold')
    ax_roc.legend(loc='lower right', frameon=True)

    ax_pr.set_xlabel("Recall", fontsize=11, fontweight='bold')
    ax_pr.set_ylabel("Precision", fontsize=11, fontweight='bold')
    ax_pr.set_title(f"Precision-Recall Curves on {dataset_name.upper()} Test Split", fontsize=12, fontweight='bold')
    ax_pr.legend(loc='lower left', frameon=True)

    plt.tight_layout()
    plot_path = os.path.join(save_dir, f"{dataset_name}_roc_pr.png")
    plt.savefig(plot_path, dpi=300)
    plt.close()
    logger.info(f"Saved ROC/PR curves -> {plot_path}")


def plot_tsne_embeddings(
    dataset_name: str,
    model_name: str = 'mpnn',
    checkpoint_dir: str = 'results/checkpoints',
    save_dir: str = 'results/plots'
):
    """Extract latent embeddings and project into 2D with t-SNE."""
    os.makedirs(save_dir, exist_ok=True)
    _, _, test_loader, meta = get_dataloaders(dataset_name, batch_size=32)

    try:
        model, _ = load_trained_model(model_name, dataset_name, checkpoint_dir)
    except Exception as e:
        logger.warning(f"Could not load {model_name} for t-SNE: {e}")
        return

    embeddings_list, targets_list = [], []
    with torch.no_grad():
        for batch in test_loader:
            emb = model.get_embedding(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch)
            embeddings_list.append(emb.numpy())
            targets_list.append(batch.y[:, 0].numpy())

    embeddings = np.concatenate(embeddings_list, axis=0)
    targets = np.concatenate(targets_list, axis=0)

    # Filter out NaNs if any
    mask = ~np.isnan(targets)
    embeddings = embeddings[mask]
    targets = targets[mask]

    # Run t-SNE
    perplexity = min(30, max(5, len(embeddings) // 4))
    tsne = TSNE(n_components=2, perplexity=perplexity, random_state=42)
    embeddings_2d = tsne.fit_transform(embeddings)

    plt.figure(figsize=(7, 6))
    if dataset_name.lower() == 'esol':
        sc = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=targets, cmap='viridis', s=50, alpha=0.85, edgecolors='k', linewidth=0.3)
        cbar = plt.colorbar(sc)
        cbar.set_label("Solubility logS (mol/L)", fontsize=10, fontweight='bold')
    else:
        # Binary classification
        scatter_0 = plt.scatter(embeddings_2d[targets == 0, 0], embeddings_2d[targets == 0, 1], c='#3b82f6', s=50, alpha=0.8, edgecolors='k', linewidth=0.3, label='Negative (0)')
        scatter_1 = plt.scatter(embeddings_2d[targets == 1, 0], embeddings_2d[targets == 1, 1], c='#ef4444', s=50, alpha=0.8, edgecolors='k', linewidth=0.3, label='Positive (1)')
        plt.legend(frameon=True)

    plt.title(f"t-SNE Latent Space ({model_name.upper()} on {dataset_name.upper()} Test Split)", fontsize=12, fontweight='bold')
    plt.xlabel("t-SNE Dimension 1", fontsize=11, fontweight='bold')
    plt.ylabel("t-SNE Dimension 2", fontsize=11, fontweight='bold')
    plt.tight_layout()

    plot_path = os.path.join(save_dir, f"{dataset_name}_tsne_{model_name}.png")
    plt.savefig(plot_path, dpi=300)
    plt.close()
    logger.info(f"Saved t-SNE visualization -> {plot_path}")


def generate_benchmark_report(
    benchmark_results: Dict[str, Dict[str, Any]],
    output_path: str = 'results/REPORT.md'
):
    """Compile comprehensive benchmark report across datasets and models."""
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    lines = [
        "# Benchmark Report: Graph Neural Networks for Molecular Property Prediction",
        "",
        "## Executive Summary",
        "This benchmark evaluates 5 distinct molecular property prediction architectures across 3 canonical MoleculeNet benchmarks using **strict Bemis-Murcko scaffold splitting (80/10/10)**.",
        "All models were trained and benchmarked locally on an Intel CPU with 4 physical threads.",
        "",
        "---",
        "",
        "## Benchmark Results",
        ""
    ]

    for dataset_name, models_dict in benchmark_results.items():
        lines.append(f"### Dataset: `{dataset_name.upper()}`")
        if dataset_name.lower() == 'esol':
            lines.append("| Model Architecture | Test RMSE (mol/L) ↓ | Test MAE (mol/L) ↓ | Test R² ↑ | Train Time (s) |")
            lines.append("| :--- | :--- | :--- | :--- | :--- |")
            for m, data in models_dict.items():
                agg = data['aggregated_test_metrics']
                rmse_str = f"{agg['rmse']['mean']:.3f} ± {agg['rmse']['std']:.3f}"
                mae_str = f"{agg['mae']['mean']:.3f} ± {agg['mae']['std']:.3f}"
                r2_str = f"{agg['r2']['mean']:.3f} ± {agg['r2']['std']:.3f}"
                time_str = f"{data['mean_train_time_sec']:.1f}s"
                lines.append(f"| **{m.upper()}** | {rmse_str} | {mae_str} | {r2_str} | {time_str} |")
        else:
            lines.append("| Model Architecture | Test ROC-AUC ↑ | Test PR-AUC ↑ | Train Time (s) |")
            lines.append("| :--- | :--- | :--- | :--- |")
            for m, data in models_dict.items():
                agg = data['aggregated_test_metrics']
                roc_str = f"{agg['roc_auc']['mean']:.3f} ± {agg['roc_auc']['std']:.3f}"
                pr_str = f"{agg['pr_auc']['mean']:.3f} ± {agg['pr_auc']['std']:.3f}"
                time_str = f"{data['mean_train_time_sec']:.1f}s"
                lines.append(f"| **{m.upper()}** | {roc_str} | {pr_str} | {time_str} |")
        lines.append("")

    lines.extend([
        "---",
        "",
        "## Inductive Biases & Architectural Insights",
        "",
        "1. **Message Passing (MPNN)**: Leverages edge attributes (bond conjugation, ring membership, stereochemistry). Performs exceptionally well on solubility where ring strain and bond conjugation directly dictate crystal lattice energy and water solvation.",
        "2. **Attention Mechanism (GAT)**: Dynamically re-weights atomic neighbor interactions through multi-head self-attention. Especially effective for recognizing functional pharmacophores without rigid topology assumptions.",
        "3. **Graph Isomorphism (GIN)**: Maximally expressive for distinguishing graph topologies via multi-layer perceptron aggregations.",
        "4. **Classical Fingerprint Baseline (Random Forest + ECFP4)**: Strong competitive baseline that serves as a sanity check. While fast, it ignores global 3D topology and continuous atom-bond interactions.",
        "",
        "---",
        "",
        "## Diagnostic Plots",
        "- **Training Trajectories**: `results/plots/{dataset}_loss_curves.png`",
        "- **Parity Plots**: `results/plots/esol_parity_{model}.png`",
        "- **ROC & PR Curves**: `results/plots/{dataset}_roc_pr.png`",
        "- **t-SNE Embeddings**: `results/plots/{dataset}_tsne_{model}.png`"
    ])

    with open(output_path, 'w') as f:
        f.write("\n".join(lines))

    logger.info(f"Report compiled successfully -> {output_path}")