#!/usr/bin/env python3
"""
DiffDock Environment Setup Checker

This script verifies that the DiffDock environment is properly configured
and all dependencies are available.

Usage:
    python setup_check.py
    python setup_check.py --verbose
"""

import argparse
import sys
import os
from pathlib import Path


def check_python_version():
    """Check Python version."""
    import sys
    version = sys.version_info

    print("Checking Python version...")
    if version.major == 3 and version.minor >= 8:
        print(f"  ✓ Python {version.major}.{version.minor}.{version.micro}")
        return True
    else:
        print(f"  ✗ Python {version.major}.{version.minor}.{version.micro} "
              f"(requires Python 3.8 or higher)")
        return False


def check_package(package_name, import_name=None, version_attr='__version__'):
    """Check if a Python package is installed."""
    if import_name is None:
        import_name = package_name

    try:
        module = __import__(import_name)
        version = getattr(module, version_attr, 'unknown')
        print(f"  ✓ {package_name:20s} (version: {version})")
        return True
    except ImportError:
        print(f"  ✗ {package_name:20s} (not installed)")
        return False


def check_pytorch():
    """Check PyTorch installation and CUDA availability."""
    print("\nChecking PyTorch...")
    try:
        import torch
        print(f"  ✓ PyTorch version: {torch.__version__}")

        # Check CUDA
        if torch.cuda.is_available():
            print(f"  ✓ CUDA available: {torch.cuda.get_device_name(0)}")
            print(f"    - CUDA version: {torch.version.cuda}")
            print(f"    - Number of GPUs: {torch.cuda.device_count()}")
            return True, True
        else:
            print(f"  ⚠ CUDA not available (will run on CPU)")
            return True, False
    except ImportError:
        print(f"  ✗ PyTorch not installed")
        return False, False


def check_pytorch_geometric():
    """Check PyTorch Geometric installation."""
    print("\nChecking PyTorch Geometric...")
    packages = [
        ('torch-geometric', 'torch_geometric'),
        ('torch-scatter', 'torch_scatter'),
        ('torch-sparse', 'torch_sparse'),
        ('torch-cluster', 'torch_cluster'),
    ]

    all_ok = True
    for pkg_name, import_name in packages:
        if not check_package(pkg_name, import_name):
            all_ok = False

    return all_ok


def check_core_dependencies():
    """Check core DiffDock dependencies."""
    print("\nChecking core dependencies...")

    dependencies = [
        ('numpy', 'numpy'),
        ('scipy', 'scipy'),
        ('pandas', 'pandas'),
        ('rdkit', 'rdkit', 'rdBase.__version__'),
        ('biopython', 'Bio', '__version__'),
        ('pytorch-lightning', 'pytorch_lightning'),
        ('PyYAML', 'yaml'),
    ]

    all_ok = True
    for dep in dependencies:
        pkg_name = dep[0]
        import_name = dep[1]
        version_attr = dep[2] if len(dep) > 2 else '__version__'

        if not check_package(pkg_name, import_name, version_attr):
            all_ok = False

    return all_ok


def check_esm():
    """Check ESM (protein language model) installation."""
    print("\nChecking ESM (for protein sequence folding)...")
    try:
        import esm
        print(f"  ✓ ESM installed (version: {esm.__version__ if hasattr(esm, '__version__') else 'unknown'})")
        return True
    except ImportError:
        print(f"  ⚠ ESM not installed (needed for protein sequence folding)")
        print(f"    Install with: pip install fair-esm")
        return False


def check_diffdock_installation():
    """Check if DiffDock is properly installed/cloned."""
    print("\nChecking DiffDock installation...")

    # Look for key files
    key_files = [
        'inference.py',
        'default_inference_args.yaml',
        'environment.yml',
    ]

    found_files = []
    missing_files = []

    for filename in key_files:
        if os.path.exists(filename):
            found_files.append(filename)
        else:
            missing_files.append(filename)

    if found_files:
        print(f"  ✓ Found DiffDock files in current directory:")
        for f in found_files:
            print(f"    - {f}")
    else:
        print(f"  ⚠ DiffDock files not found in current directory")
        print(f"    Current directory: {os.getcwd()}")
        print(f"    Make sure you're in the DiffDock repository root")

    # Check for model checkpoints
    model_dir = Path('./workdir/v1.1/score_model')
    confidence_dir = Path('./workdir/v1.1/confidence_model')

    if model_dir.exists() and confidence_dir.exists():
        print(f"  ✓ Model checkpoints found")
    else:
        print(f"  ⚠ Model checkpoints not found in ./workdir/v1.1/")
        print(f"    Models will be downloaded on first run")

    return len(found_files) > 0


def print_installation_instructions():
    """Print installation instructions if setup is incomplete."""
    print("\n" + "="*80)
    print("Installation Instructions")
    print("="*80)

    print("""
If DiffDock is not installed, follow these steps:

1. Clone the repository:
   git clone https://github.com/gcorso/DiffDock.git
   cd DiffDock

2. Create conda environment:
   conda env create --file environment.yml
   conda activate diffdock

3. Verify installation:
   python setup_check.py

For Docker installation:
   docker pull rbgcsail/diffdock
   docker run -it --gpus all --entrypoint /bin/bash rbgcsail/diffdock
   micromamba activate diffdock

For more information, visit: https://github.com/gcorso/DiffDock
    """)


def print_performance_notes(has_cuda):
    """Print performance notes based on available hardware."""
    print("\n" + "="*80)
    print("Performance Notes")
    print("="*80)

    if has_cuda:
        print("""
✓ GPU detected - DiffDock will run efficiently

Expected performance:
  - First run: ~2-5 minutes (pre-computing SO(2)/SO(3) tables)
  - Subsequent runs: ~10-60 seconds per complex (depending on settings)
  - Batch processing: Highly efficient with GPU
        """)
    else:
        print("""
⚠ No GPU detected - DiffDock will run on CPU

Expected performance:
  - CPU inference is SIGNIFICANTLY slower than GPU
  - Single complex: Several minutes to hours
  - Batch processing: Not recommended on CPU

Recommendation: Use GPU for practical applications
  - Cloud options: Google Colab, AWS, or other cloud GPU services
  - Local: Install CUDA-capable GPU
        """)


def main():
    parser = argparse.ArgumentParser(
        description='Check DiffDock environment setup',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument('--verbose', '-v', action='store_true',
                        help='Show detailed version information')

    args = parser.parse_args()

    print("="*80)
    print("DiffDock Environment Setup Checker")
    print("="*80)

    checks = []

    # Run all checks
    checks.append(("Python version", check_python_version()))

    pytorch_ok, has_cuda = check_pytorch()
    checks.append(("PyTorch", pytorch_ok))

    checks.append(("PyTorch Geometric", check_pytorch_geometric()))
    checks.append(("Core dependencies", check_core_dependencies()))
    checks.append(("ESM", check_esm()))
    checks.append(("DiffDock files", check_diffdock_installation()))

    # Summary
    print("\n" + "="*80)
    print("Summary")
    print("="*80)

    all_passed = all(result for _, result in checks)

    for check_name, result in checks:
        status = "✓ PASS" if result else "✗ FAIL"
        print(f"  {status:8s} - {check_name}")

    if all_passed:
        print("\n✓ All checks passed! DiffDock is ready to use.")
        print_performance_notes(has_cuda)
        return 0
    else:
        print("\n✗ Some checks failed. Please install missing dependencies.")
        print_installation_instructions()
        return 1


if __name__ == '__main__':
    sys.exit(main())