#!/usr/bin/env python3
"""
TDC Molecular Generation with Oracles Template

This script demonstrates how to use TDC oracles for molecular generation
tasks including goal-directed generation and distribution learning.

Usage:
    python molecular_generation.py
"""

from tdc.generation import MolGen
from tdc import Oracle
import numpy as np


def load_generation_dataset():
    """
    Load molecular generation dataset
    """
    print("=" * 60)
    print("Loading Molecular Generation Dataset")
    print("=" * 60)

    # Load ChEMBL dataset
    data = MolGen(name='ChEMBL_V29')

    # Get training molecules
    split = data.get_split()
    train_smiles = split['train']['Drug'].tolist()

    print(f"\nDataset: ChEMBL_V29")
    print(f"Training molecules: {len(train_smiles)}")

    # Display sample molecules
    print("\nSample SMILES:")
    for i, smiles in enumerate(train_smiles[:5], 1):
        print(f"  {i}. {smiles}")

    return train_smiles


def single_oracle_example():
    """
    Example: Using a single oracle for molecular evaluation
    """
    print("\n" + "=" * 60)
    print("Example 1: Single Oracle Evaluation")
    print("=" * 60)

    # Initialize oracle for GSK3B target
    oracle = Oracle(name='GSK3B')

    # Test molecules
    test_molecules = [
        'CC(C)Cc1ccc(cc1)C(C)C(O)=O',  # Ibuprofen
        'CC(=O)Oc1ccccc1C(=O)O',        # Aspirin
        'Cn1c(=O)c2c(ncn2C)n(C)c1=O',   # Caffeine
        'CN1C=NC2=C1C(=O)N(C(=O)N2C)C'  # Theophylline
    ]

    print("\nEvaluating molecules with GSK3B oracle:")
    print("-" * 60)

    for smiles in test_molecules:
        score = oracle(smiles)
        print(f"SMILES: {smiles}")
        print(f"GSK3B score: {score:.4f}\n")


def multiple_oracles_example():
    """
    Example: Using multiple oracles for multi-objective optimization
    """
    print("\n" + "=" * 60)
    print("Example 2: Multiple Oracles (Multi-Objective)")
    print("=" * 60)

    # Initialize multiple oracles
    oracles = {
        'QED': Oracle(name='QED'),        # Drug-likeness
        'SA': Oracle(name='SA'),          # Synthetic accessibility
        'GSK3B': Oracle(name='GSK3B'),    # Target binding
        'LogP': Oracle(name='LogP')       # Lipophilicity
    }

    # Test molecule
    test_smiles = 'CC(C)Cc1ccc(cc1)C(C)C(O)=O'

    print(f"\nEvaluating: {test_smiles}")
    print("-" * 60)

    scores = {}
    for name, oracle in oracles.items():
        score = oracle(test_smiles)
        scores[name] = score
        print(f"{name:10s}: {score:.4f}")

    # Multi-objective score (weighted combination)
    print("\n--- Multi-Objective Scoring ---")

    # Invert SA (lower is better, so we invert for maximization)
    sa_score = 1.0 / (1.0 + scores['SA'])

    # Weighted combination
    weights = {'QED': 0.3, 'SA': 0.2, 'GSK3B': 0.4, 'LogP': 0.1}
    multi_score = (
        weights['QED'] * scores['QED'] +
        weights['SA'] * sa_score +
        weights['GSK3B'] * scores['GSK3B'] +
        weights['LogP'] * (scores['LogP'] / 5.0)  # Normalize LogP
    )

    print(f"Multi-objective score: {multi_score:.4f}")
    print(f"Weights: {weights}")


def batch_evaluation_example():
    """
    Example: Batch evaluation of multiple molecules
    """
    print("\n" + "=" * 60)
    print("Example 3: Batch Evaluation")
    print("=" * 60)

    # Generate sample molecules
    molecules = [
        'CC(C)Cc1ccc(cc1)C(C)C(O)=O',
        'CC(=O)Oc1ccccc1C(=O)O',
        'Cn1c(=O)c2c(ncn2C)n(C)c1=O',
        'CN1C=NC2=C1C(=O)N(C(=O)N2C)C',
        'CC(C)NCC(COc1ccc(cc1)COCCOC(C)C)O'
    ]

    # Initialize oracle
    oracle = Oracle(name='DRD2')

    print(f"\nBatch evaluating {len(molecules)} molecules with DRD2 oracle...")

    # Batch evaluation (more efficient than individual calls)
    scores = oracle(molecules)

    print("\nResults:")
    print("-" * 60)
    for smiles, score in zip(molecules, scores):
        print(f"{smiles[:40]:40s}... Score: {score:.4f}")

    # Statistics
    print(f"\nStatistics:")
    print(f"  Mean score: {np.mean(scores):.4f}")
    print(f"  Std score: {np.std(scores):.4f}")
    print(f"  Min score: {np.min(scores):.4f}")
    print(f"  Max score: {np.max(scores):.4f}")


def goal_directed_generation_template():
    """
    Template for goal-directed molecular generation
    """
    print("\n" + "=" * 60)
    print("Example 4: Goal-Directed Generation Template")
    print("=" * 60)

    template = '''
# Template for goal-directed molecular generation

from tdc.generation import MolGen
from tdc import Oracle
import numpy as np

# 1. Load training data
data = MolGen(name='ChEMBL_V29')
train_smiles = data.get_split()['train']['Drug'].tolist()

# 2. Initialize oracle(s)
oracle = Oracle(name='GSK3B')

# 3. Initialize your generative model
# model = YourGenerativeModel()
# model.fit(train_smiles)

# 4. Generation loop
num_iterations = 100
num_molecules_per_iter = 100
best_molecules = []

for iteration in range(num_iterations):
    # Generate candidate molecules
    # candidates = model.generate(num_molecules_per_iter)

    # Evaluate with oracle
    scores = oracle(candidates)

    # Select top molecules
    top_indices = np.argsort(scores)[-10:]
    top_molecules = [candidates[i] for i in top_indices]
    top_scores = [scores[i] for i in top_indices]

    # Store best molecules
    best_molecules.extend(zip(top_molecules, top_scores))

    # Optional: Fine-tune model on top molecules
    # model.fine_tune(top_molecules)

    # Print progress
    print(f"Iteration {iteration}: Best score = {max(scores):.4f}")

# Sort and display top molecules
best_molecules.sort(key=lambda x: x[1], reverse=True)
print("\\nTop 10 molecules:")
for smiles, score in best_molecules[:10]:
    print(f"{smiles}: {score:.4f}")
'''

    print("\nGoal-Directed Generation Template:")
    print("=" * 60)
    print(template)


def distribution_learning_example(train_smiles):
    """
    Example: Distribution learning evaluation
    """
    print("\n" + "=" * 60)
    print("Example 5: Distribution Learning")
    print("=" * 60)

    # Use subset for demonstration
    train_subset = train_smiles[:1000]

    # Initialize oracle
    oracle = Oracle(name='QED')

    print("\nEvaluating property distribution...")

    # Evaluate training set
    print("Computing training set distribution...")
    train_scores = oracle(train_subset)

    # Simulate generated molecules (in practice, use your generative model)
    # For demo: add noise to training molecules
    print("Computing generated set distribution...")
    generated_scores = train_scores + np.random.normal(0, 0.1, len(train_scores))
    generated_scores = np.clip(generated_scores, 0, 1)  # QED is [0, 1]

    # Compare distributions
    print("\n--- Distribution Statistics ---")
    print(f"Training set (n={len(train_subset)}):")
    print(f"  Mean: {np.mean(train_scores):.4f}")
    print(f"  Std: {np.std(train_scores):.4f}")
    print(f"  Median: {np.median(train_scores):.4f}")

    print(f"\nGenerated set (n={len(generated_scores)}):")
    print(f"  Mean: {np.mean(generated_scores):.4f}")
    print(f"  Std: {np.std(generated_scores):.4f}")
    print(f"  Median: {np.median(generated_scores):.4f}")

    # Distribution similarity metrics
    from scipy.stats import ks_2samp
    ks_statistic, p_value = ks_2samp(train_scores, generated_scores)

    print(f"\nKolmogorov-Smirnov Test:")
    print(f"  KS statistic: {ks_statistic:.4f}")
    print(f"  P-value: {p_value:.4f}")

    if p_value > 0.05:
        print("  → Distributions are similar (p > 0.05)")
    else:
        print("  → Distributions are significantly different (p < 0.05)")


def available_oracles_info():
    """
    Display information about available oracles
    """
    print("\n" + "=" * 60)
    print("Example 6: Available Oracles")
    print("=" * 60)

    oracle_info = {
        'Biochemical Targets': [
            'DRD2', 'GSK3B', 'JNK3', '5HT2A', 'ACE',
            'MAPK', 'CDK', 'P38', 'PARP1', 'PIK3CA'
        ],
        'Physicochemical Properties': [
            'QED', 'SA', 'LogP', 'MW', 'Lipinski'
        ],
        'Composite Metrics': [
            'Isomer_Meta', 'Median1', 'Median2',
            'Rediscovery', 'Similarity', 'Uniqueness', 'Novelty'
        ],
        'Specialized': [
            'ASKCOS', 'Docking', 'Vina'
        ]
    }

    print("\nAvailable Oracle Categories:")
    print("-" * 60)

    for category, oracles in oracle_info.items():
        print(f"\n{category}:")
        for oracle_name in oracles:
            print(f"  - {oracle_name}")

    print("\nFor detailed oracle documentation, see:")
    print("  references/oracles.md")


def constraint_satisfaction_example():
    """
    Example: Molecular generation with constraints
    """
    print("\n" + "=" * 60)
    print("Example 7: Constraint Satisfaction")
    print("=" * 60)

    # Define constraints
    constraints = {
        'QED': (0.5, 1.0),      # Drug-likeness >= 0.5
        'SA': (1.0, 5.0),       # Easy to synthesize
        'MW': (200, 500),       # Molecular weight 200-500 Da
        'LogP': (0, 3)          # Lipophilicity 0-3
    }

    # Initialize oracles
    oracles = {name: Oracle(name=name) for name in constraints.keys()}

    # Test molecules
    test_molecules = [
        'CC(C)Cc1ccc(cc1)C(C)C(O)=O',
        'CC(=O)Oc1ccccc1C(=O)O',
        'Cn1c(=O)c2c(ncn2C)n(C)c1=O'
    ]

    print("\nConstraints:")
    for prop, (min_val, max_val) in constraints.items():
        print(f"  {prop}: [{min_val}, {max_val}]")

    print("\n" + "-" * 60)
    print("Evaluating molecules against constraints:")
    print("-" * 60)

    for smiles in test_molecules:
        print(f"\nSMILES: {smiles}")

        satisfies_all = True
        for prop, (min_val, max_val) in constraints.items():
            score = oracles[prop](smiles)
            satisfies = min_val <= score <= max_val

            status = "✓" if satisfies else "✗"
            print(f"  {prop:10s}: {score:7.2f} [{min_val:5.1f}, {max_val:5.1f}] {status}")

            satisfies_all = satisfies_all and satisfies

        result = "PASS" if satisfies_all else "FAIL"
        print(f"  Overall: {result}")


def main():
    """
    Main function to run all molecular generation examples
    """
    print("\n" + "=" * 60)
    print("TDC Molecular Generation with Oracles Examples")
    print("=" * 60)

    # Load generation dataset
    train_smiles = load_generation_dataset()

    # Example 1: Single oracle
    single_oracle_example()

    # Example 2: Multiple oracles
    multiple_oracles_example()

    # Example 3: Batch evaluation
    batch_evaluation_example()

    # Example 4: Goal-directed generation template
    goal_directed_generation_template()

    # Example 5: Distribution learning
    distribution_learning_example(train_smiles)

    # Example 6: Available oracles
    available_oracles_info()

    # Example 7: Constraint satisfaction
    constraint_satisfaction_example()

    print("\n" + "=" * 60)
    print("Molecular generation examples completed!")
    print("=" * 60)
    print("\nNext steps:")
    print("1. Implement your generative model")
    print("2. Use oracles to guide generation")
    print("3. Evaluate generated molecules")
    print("4. Iterate and optimize")
    print("=" * 60)


if __name__ == "__main__":
    main()