"""
Main entry script for PandaDock with GPU/CPU hardware acceleration.
"""

import argparse
import os
import time
import sys
import json
from datetime import datetime
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import traceback
from .protein import Protein
from .ligand import Ligand
from .scoring import CompositeScoringFunction, EnhancedScoringFunction
from .search import RandomSearch, GeneticAlgorithm
from .utils import (
    setup_logging,
    save_docking_results,
    create_initial_files,
    update_status,
    save_intermediate_result,
    create_descriptive_output_dir
)
from .preparation import prepare_protein, prepare_ligand
from .reporting import DockingReporter
from .validation import validate_against_reference
from .main_integration import (
    add_hardware_options, 
    configure_hardware,
    setup_hardware_acceleration,
    create_optimized_scoring_function,
    create_optimized_search_algorithm,
    get_scoring_type_from_args,
    get_algorithm_type_from_args,
    get_algorithm_kwargs_from_args
)

from . import __version__

__all__ = ['__version__', 'add_hardware_options', 'configure_hardware', 
           'setup_hardware_acceleration', 'create_optimized_scoring_function', 
           'create_optimized_search_algorithm', 'get_scoring_type_from_args', 
           'get_algorithm_type_from_args', 'get_algorithm_kwargs_from_args']

# Import physics-based algorithms
try:
    from .physics import (MMFFMinimization, GeneralizedBornSolvation, 
                          MonteCarloSampling, PhysicsBasedScoring)
    PHYSICS_AVAILABLE = True
except ImportError:
    PHYSICS_AVAILABLE = False
    print("Warning: Physics-based modules not available. Some features will be disabled.")

def check_for_updates(logger):
    """Check for newer versions of PandaDock on PyPI and notify user if available."""
    try:
        import os
        import time
        import requests
        import pkg_resources
        from packaging import version

        # Check cache - only verify once per day
        cache_file = os.path.join(os.path.expanduser("~"), ".pandadock_version_check")
        current_time = time.time()
        
        # If cache exists and is less than 24 hours old, skip check
        if os.path.exists(cache_file):
            with open(cache_file, "r") as f:
                last_check = float(f.read().strip())
            if current_time - last_check < 86400:  # 24 hours
                return
            
        # Update cache timestamp
        with open(cache_file, "w") as f:
            f.write(str(current_time))
            
        # Get current installed version
        current_version = pkg_resources.get_distribution("pandadock").version
        
        # Query PyPI for the latest version
        response = requests.get("https://pypi.org/pypi/pandadock/json", timeout=2)
        latest_version = response.json()["info"]["version"]
        
        # Compare versions
        if version.parse(latest_version) > version.parse(current_version):
            logger.info("\n" + "="*70)
            logger.info(f"  New version available: PandaDock {latest_version} (you have {current_version})")
            logger.info(f"  Update with: pip install --upgrade pandadock")
            logger.info("  See release notes at: https://github.com/pritampanda15/PandaDock/releases")
            logger.info("="*70 + "\n")
            
    except Exception:
        # Silently fail if version check doesn't work
        pass

def prepare_protein_configs(protein, args, logger):
    """
    Prepare protein configurations for benchmarking.
    
    Parameters:
    -----------
    protein : Protein
        Protein object
    args : argparse.Namespace
        Command-line arguments
    logger : logging.Logger
        Logger instance
    
    Returns:
    --------
    list
        List of protein configurations for benchmarking
    """
    configs = []
    
    # Rigid configuration
    rigid_protein = protein  # No need to copy as we're creating new flex_protein if needed
    configs.append({
        'type': 'rigid',
        'protein': rigid_protein,
        'flexible_residues': []
    })
    
    # Flexible configuration if requested
    if hasattr(args, 'flex_residues') and args.flex_residues or (hasattr(args, 'auto_flex') and args.auto_flex):
        import copy
        flex_protein = copy.deepcopy(protein)
        flex_residues = []
        
        if hasattr(args, 'flex_residues') and args.flex_residues:
            # Use user-specified flexible residues
            flex_residues = args.flex_residues
            logger.info(f"Using user-specified flexible residues: {', '.join(flex_residues)}")
        elif hasattr(args, 'auto_flex') and args.auto_flex:
            # Auto-detect flexible residues based on binding site
            if protein.active_site and 'residues' in protein.active_site:
                # Get residues in the active site
                binding_site_residues = protein.active_site['residues']
                
                # This function needs to be added to the appropriate class
                flex_residues = _detect_flexible_residues(protein, binding_site_residues, 
                                                          max_residues=args.max_flex_residues if hasattr(args, 'max_flex_residues') else 5,
                                                          logger=logger)
                
                logger.info(f"Auto-detected flexible residues: {', '.join(flex_residues)}")
            else:
                logger.info("Warning: No active site defined. Cannot auto-detect flexible residues.")
        
        if flex_residues:
            flex_protein.define_flexible_residues(flex_residues, 
                                                 max_rotatable_bonds=args.max_flex_bonds if hasattr(args, 'max_flex_bonds') else 3)
            configs.append({
                'type': 'flexible',
                'protein': flex_protein,
                'flexible_residues': flex_residues
            })
        else:
            logger.info("No flexible residues defined. Using only rigid configuration.")
    
    return configs

def _detect_flexible_residues(protein, binding_site_residues, max_residues=5, logger=None):
    """
    Automatically detect flexible residues in the binding site.
    
    Parameters:
    -----------
    protein : Protein
        Protein object
    binding_site_residues : list
        List of residue IDs in the binding site
    max_residues : int
        Maximum number of flexible residues to detect
    logger : logging.Logger
        Logger instance
    
    Returns:
    --------
    list
        List of detected flexible residue IDs
    """
    # Define residues with flexible sidechains
    flexible_aa_types = [
        'ARG',  # Arginine - very flexible sidechain with multiple rotatable bonds
        'LYS',  # Lysine - long flexible sidechain
        'GLU',  # Glutamic acid - flexible sidechain with carboxyl group
        'GLN',  # Glutamine - flexible sidechain with amide group
        'MET',  # Methionine - flexible sidechain with sulfur
        'PHE',  # Phenylalanine - aromatic sidechain that can rotate
        'TYR',  # Tyrosine - aromatic sidechain with hydroxyl group
        'TRP',  # Tryptophan - large aromatic sidechain
        'LEU',  # Leucine - branched hydrophobic sidechain
        'ILE',  # Isoleucine - branched hydrophobic sidechain
        'ASP',  # Aspartic acid - shorter version of glutamic acid
        'ASN',  # Asparagine - shorter version of glutamine
        'HIS',  # Histidine - aromatic sidechain that can rotate
        'SER',  # Serine - small sidechain with hydroxyl group
        'THR',  # Threonine - small sidechain with hydroxyl group
        'CYS',  # Cysteine - sidechain with thiol group
        'VAL'   # Valine - smaller branched hydrophobic sidechain
    ]
    
    if logger:
        logger.info(f"Searching for flexible residues among {len(binding_site_residues)} binding site residues")
    
    candidate_residues = []
    
    for res_id in binding_site_residues:
        if res_id in protein.residues:
            residue_atoms = protein.residues[res_id]
            
            # Get residue type
            if residue_atoms and 'residue_name' in residue_atoms[0]:
                res_type = residue_atoms[0]['residue_name']
                
                # Check if it's a residue type with flexible sidechain
                if res_type in flexible_aa_types:
                    # Calculate distance from binding site center
                    if protein.active_site and 'center' in protein.active_site:
                        center = protein.active_site['center']
                        
                        # Use CA atom or any atom for distance calculation
                        ca_atom = next((atom for atom in residue_atoms if atom.get('name', '') == 'CA'), None)
                        
                        if ca_atom:
                            distance = np.linalg.norm(ca_atom['coords'] - center)
                            candidate_residues.append((res_id, distance, res_type))
                            if logger:
                                logger.info(f"  Found candidate flexible residue: {res_id} ({res_type}) - distance: {distance:.2f}Å")
    
    # Sort by distance to center (closest first)
    candidate_residues.sort(key=lambda x: x[1])
    
    if logger:
        logger.info(f"Selected {min(max_residues, len(candidate_residues))} flexible residues:")
        for i, (res_id, distance, res_type) in enumerate(candidate_residues[:max_residues]):
            logger.info(f"  {i+1}. {res_id} ({res_type}) - distance: {distance:.2f}Å")
    
    # Return up to max_residues
    return [res_id for res_id, _, _ in candidate_residues[:max_residues]]

def write_results_to_txt(results, output_dir, elapsed_time, protein_path, ligand_path, algorithm, iterations, logger):
    """
    Write detailed docking results to a text file.
    
    Parameters:
    -----------
    results : list
        List of (pose, score) tuples
    output_dir : str
        Output directory
    elapsed_time : float
        Total elapsed time in seconds
    protein_path : str
        Path to protein file
    ligand_path : str
        Path to ligand file
    algorithm : str
        Docking algorithm used
    iterations : int
        Number of iterations/generations
    logger : logging.Logger
        Logger instance
    """
    results_path = Path(output_dir) / "docking_results.txt"
    
    with open(results_path, 'w') as f:
        f.write("=====================================================\n")
        f.write("        PandaDock - Python Molecular Docking Results    \n")
        f.write("=====================================================\n\n")
        
        # Write run information
        f.write("RUN INFORMATION\n")
        f.write("--------------\n")
        f.write(f"Date and Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Protein: {protein_path}\n")
        f.write(f"Ligand: {ligand_path}\n")
        f.write(f"Algorithm: {algorithm}\n")
        f.write(f"Iterations/Generations: {iterations}\n")
        f.write(f"Total Runtime: {elapsed_time:.2f} seconds\n\n")
        
        # Check if results is empty
        if not results:
            f.write("RESULTS SUMMARY\n")
            f.write("--------------\n")
            f.write("No valid docking solutions found.\n")
            f.write("This can occur due to incompatible structures, overly strict scoring parameters,\n")
            f.write("or issues with the search space definition.\n\n")
        else:
            # Sort results by score (lowest first)
            sorted_results = sorted(results, key=lambda x: x[1])
            
            # Write summary of results
            f.write("RESULTS SUMMARY\n")
            f.write("--------------\n")
            f.write(f"Total Poses Generated: {len(results)}\n")
            f.write(f"Best Score: {sorted_results[0][1]:.4f}\n")
            f.write(f"Worst Score: {sorted_results[-1][1]:.4f}\n")
            f.write(f"Average Score: {sum([score for _, score in results])/len(results):.4f}\n\n")
            
            # Write top 10 poses
            f.write("TOP 10 POSES\n")
            f.write("--------------\n")
            f.write("Rank\tScore\tFile\n")
            for i, (pose, score) in enumerate(sorted_results[:10]):
                f.write(f"{i+1}\t{score:.4f}\tpose_{i+1}_score_{score:.1f}.pdb\n")
            
            f.write("\n\nFull results are available in the output directory.\n")
        
        f.write("=====================================================\n")
    
    if logger:
        logger.info(f"Detailed results written to {results_path}")

def add_advanced_search_options(parser):
    """Add command-line options for advanced search algorithms."""
    adv_search = parser.add_argument_group('Advanced Search Algorithms')
    
    # Algorithm selection
    adv_search.add_argument('--advanced-search', choices=['gradient', 'replica-exchange', 
                                                         'ml-guided', 'fragment-based', 'hybrid'],
                           help='Advanced search algorithm to use')
    
    # Gradient-based options
    adv_search.add_argument('--gradient-step', type=float, default=0.1,
                          help='Step size for gradient calculation in gradient-based search')
    adv_search.add_argument('--convergence-threshold', type=float, default=0.01,
                          help='Convergence threshold for gradient-based search')
    
    # Replica exchange options
    adv_search.add_argument('--n-replicas', type=int, default=4,
                          help='Number of replicas for replica exchange')
    adv_search.add_argument('--replica-temperatures', type=float, nargs='+',
                          help='Temperatures for replicas (e.g., 300 400 500 600)')
    adv_search.add_argument('--exchange-steps', type=int, default=10,
                          help='Number of exchange attempts in replica exchange')
    
    # ML-guided options
    adv_search.add_argument('--surrogate-model', choices=['rf', 'gp', 'nn'], default='rf',
                          help='Surrogate model type for ML-guided search')
    adv_search.add_argument('--exploitation-factor', type=float, default=0.8,
                          help='Exploitation vs exploration balance (0-1) for ML-guided search')
    
    # Fragment-based options
    adv_search.add_argument('--fragment-min-size', type=int, default=5,
                          help='Minimum fragment size for fragment-based docking')
    adv_search.add_argument('--growth-steps', type=int, default=3,
                          help='Number of fragment growth steps')
    
    # Hybrid search options
    adv_search.add_argument('--ga-iterations', type=int, default=50,
                          help='Genetic algorithm iterations in hybrid search')
    adv_search.add_argument('--lbfgs-iterations', type=int, default=50,
                          help='L-BFGS iterations in hybrid search')
    adv_search.add_argument('--top-n-for-local', type=int, default=10,
                          help='Top N poses to optimize with L-BFGS in hybrid search')

def add_analysis_options(parser):
    """Add command-line options for pose clustering and analysis."""
    analysis = parser.add_argument_group('Pose Clustering and Analysis')
    
    # Clustering options
    analysis.add_argument('--cluster-poses', action='store_true',
                         help='Perform clustering of docking poses')
    analysis.add_argument('--clustering-method', choices=['hierarchical', 'dbscan'], 
                        default='hierarchical',
                        help='Method for clustering poses')
    analysis.add_argument('--rmsd-cutoff', type=float, default=2.0,
                        help='RMSD cutoff for pose clustering')
    
    # Interaction analysis
    analysis.add_argument('--analyze-interactions', action='store_true',
                         help='Generate interaction fingerprints and analysis')
    analysis.add_argument('--interaction-types', nargs='+',
                        choices=['hbond', 'hydrophobic', 'ionic', 'aromatic', 'halogen'],
                        default=['hbond', 'hydrophobic', 'ionic'],
                        help='Interaction types to include in analysis')
    
    # Binding mode analysis
    analysis.add_argument('--classify-modes', action='store_true',
                         help='Classify binding modes of docking poses')
    analysis.add_argument('--discover-modes', action='store_true',
                         help='Automatically discover binding modes from results')
    analysis.add_argument('--n-modes', type=int, default=5,
                        help='Number of binding modes to discover')
    
    # Energy analysis
    analysis.add_argument('--energy-decomposition', action='store_true',
                         help='Perform energy decomposition analysis')
    analysis.add_argument('--per-residue-energy', action='store_true',
                         help='Calculate per-residue energy contributions')
    
    # Reporting - with renamed arguments to avoid conflicts
    analysis.add_argument('--generate-analysis-report', action='store_true',
                         help='Generate comprehensive docking report')
    analysis.add_argument('--analysis-report-format', choices=['html', 'pdf', 'txt'], default='html',
                        help='Format for analysis report')
    analysis.add_argument('--analysis-report-sections', nargs='+',
                        choices=['summary', 'clusters', 'interactions', 'energetics'],
                        default=['summary', 'clusters', 'interactions', 'energetics'],
                        help='Sections to include in the analysis report')

def main():
    # Initialize return value
    return_code = 0

    # Record start time immediately
    start_time = time.time()

    # Initialize variables that might be used in finally block
    temp_dir = None
    hybrid_manager = None
    output_dir = None
    
    try:
        # Parse command line arguments
        parser = argparse.ArgumentParser(description='PandaDock: Python Molecular Docking Tool')
        
        # Required arguments
        parser.add_argument('-p', '--protein', required=True, help='Path to protein PDB file')
        parser.add_argument('-l', '--ligand', required=True, help='Path to ligand MOL/SDF file')
        
        # Optional arguments
        parser.add_argument('-o', '--output', default='docking_results', 
                            help='Output directory for docking results')
        parser.add_argument('-a', '--algorithm', choices=['random', 'genetic', 'pandadock'], default='genetic',
                            help='Docking algorithm to use (default: genetic)')
        parser.add_argument('-i', '--iterations', type=int, default=1000,
                            help='Number of iterations/generations (default: 1000)')
        parser.add_argument('-s', '--site', nargs=3, type=float, metavar=('X', 'Y', 'Z'),
                            help='Active site center coordinates')
        parser.add_argument('-r', '--radius', type=float, default=10.0,
                            help='Active site radius in Angstroms (default: 10.0)')
        parser.add_argument('--detect-pockets', action='store_true',
                            help='Automatically detect binding pockets')
        
        parser.add_argument('--tethered-docking', action='store_true',
                    help='Use tethered scoring with reference structure')
        parser.add_argument('--tether-weight', type=float, default=10.0,
                    help='Weight for tethered scoring (higher = stronger tethering)')
        
        # Quick mode options
        parser.add_argument('--fast-mode', action='store_true',
                            help='Run with minimal enhancements for quick results')
        parser.add_argument('--enhanced', action='store_true',
                            help='Use enhanced algorithms for more accurate (but slower) results')
        
        # Enhanced docking options
        parser.add_argument('--enhanced-scoring', action='store_true',
                            help='Use enhanced scoring function with electrostatics')
        parser.add_argument('--prepare-molecules', action='store_true',
                            help='Prepare protein and ligand before docking (recommended)')
        parser.add_argument('--reference', 
                            help='Reference ligand structure for validation')
        parser.add_argument('--exact-alignment', action='store_true',
                        help='Align docked pose exactly to reference structure')
        parser.add_argument('--population-size', type=int, default=100,
                            help='Population size for genetic algorithm (default: 100)')
        parser.add_argument('--exhaustiveness', type=int, default=1,
                            help='Number of independent docking runs (default: 1)')
        parser.add_argument('--local-opt', action='store_true', # Default is False
                            help='Enable local optimization on top poses (default: disabled)')
        parser.add_argument('--ph', type=float, default=7.4,
                            help='pH for protein preparation (default: 7.4)')
        
        # Physics-based options
        parser.add_argument('--physics-based', action='store_true',
                            help='Use full physics-based scoring (very slow but most accurate)')
        parser.add_argument('--mmff-minimization', action='store_true',
                            help='Use MMFF94 force field minimization (requires RDKit)')
        parser.add_argument('--monte-carlo', action='store_true',
                            help='Use Monte Carlo sampling instead of genetic algorithm')
        parser.add_argument('--mc-steps', type=int, default=1000,
                            help='Number of Monte Carlo steps (default: 1000)')
        parser.add_argument('--temperature', type=float, default=300.0,
                            help='Temperature for Monte Carlo simulation in Kelvin (default: 300K)')

        # PandaDock options
        pandadock_group = parser.add_argument_group('PandaDock Options')
        pandadock_group.add_argument('--high-temp', type=float, default=1000.0,
                                help='High temperature for pandadock MD simulations (K)')
        pandadock_group.add_argument('--target-temp', type=float, default=300.0,
                                help='Target temperature for pandadock cooling (K)')
        pandadock_group.add_argument('--num-conformers', type=int, default=10,
                                help='Number of ligand conformers to generate in pandadock')
        pandadock_group.add_argument('--num-orientations', type=int, default=10,
                                help='Number of orientations to try for each conformer in pandadock')
        pandadock_group.add_argument('--md-steps', type=int, default=1000,
                                help='Number of MD steps for simulated annealing in pandadock')
        pandadock_group.add_argument('--minimize-steps', type=int, default=200,
                                help='Number of minimization steps for final refinement in pandadock')
        pandadock_group.add_argument('--use-grid', action='store_true',
                                help='Use grid-based energy calculations in pandadock')
        pandadock_group.add_argument('--cooling-factor', type=float, default=0.95,
                                help='Cooling factor for simulated annealing (applies to PANDADOCK and Monte Carlo)')

        # Auto-algorithm option
        parser.add_argument('--auto-algorithm', action='store_true',
                        help='Automatically select the best docking algorithm based on your system')

        # Flexible residue options
        parser.add_argument('--flex-residues', nargs='+', 
                            help='Specify flexible residue IDs (e.g., A_42 B_57)')
        parser.add_argument('--max-flex-bonds', type=int, default=3,
                            help='Maximum rotatable bonds per residue (default: 3)')
        
        parser.add_argument('--auto-flex', action='store_true',
                        help='Automatically detect flexible residues in the binding site')
        parser.add_argument('--max-flex-residues', type=int, default=5,
                        help='Maximum number of flexible residues to detect in auto mode (default: 5)')
        
        # Add reporting options
        report_group = parser.add_argument_group('Reporting Options')
        report_group.add_argument('--report-format', choices=['text', 'csv', 'json', 'html', 'all'],
                                default='all', help='Report format (default: all)')
        report_group.add_argument('--report-name', type=str, default=None,
                                help='Custom name for the report files')
        report_group.add_argument('--detailed-energy', action='store_true',
                                help='Include detailed energy component breakdown in reports')
        report_group.add_argument('--skip-plots', action='store_true',
                                help='Skip generating plots for reports')
        
        # Add hardware acceleration options
        add_hardware_options(parser)
        add_advanced_search_options(parser)
        add_analysis_options(parser)
        
        args = parser.parse_args()
        
        # Create descriptive output directory name
        protein_base = Path(args.protein).stem
        ligand_base = Path(args.ligand).stem
        algo_name = args.algorithm
        readable_date = datetime.now().strftime("%Y-%m-%d_%H-%M")
        output_name = f"{protein_base}_{ligand_base}_{algo_name}_{readable_date}"
        
        # Setup output directory
        output_dir = f"{args.output}_{output_name}"
        os.makedirs(output_dir, exist_ok=True)
        
        # Create initial files and setup logging
        create_initial_files(output_dir, args)
        logger = setup_logging(output_dir)
        logger.info(f"============ PandaDock - Python Molecular Docking ============")
        logger.info(f"PandaDock starting - output will be saved to {output_dir}")
        
        # Check for updates at startup
        check_for_updates(logger)
        
        # Configure hardware settings
        hw_config = configure_hardware(args)
        
        # Create temporary directory for prepared files
        temp_dir = Path('temp_pandadock')
        os.makedirs(temp_dir, exist_ok=True)

        # Process mode flags - set parameters based on mode
        if args.auto_algorithm:
            # Simple heuristic for algorithm selection:
            # - If physics-based scoring is enabled, use PANDADOCK
            # - If enhanced scoring is enabled with local opt, use Monte Carlo
            # - If fast mode is enabled, use Random
            # - Otherwise, use Genetic
            if args.physics_based:
                args.algorithm = 'pandadock'
                logger.info("Auto-selecting PANDADOCK algorithm for physics-based scoring")
            elif args.enhanced_scoring and args.local_opt:
                args.monte_carlo = True  # This overrides args.algorithm
                logger.info("Auto-selecting Monte Carlo algorithm for enhanced scoring with local optimization")
            elif args.fast_mode:
                args.algorithm = 'random'
                logger.info("Auto-selecting Random search algorithm for fast mode")
            else:
                args.algorithm = 'genetic'
                logger.info("Auto-selecting Genetic algorithm (default)")
      
        if args.fast_mode:
            logger.info("\nRunning in fast mode with minimal enhancements")
            args.enhanced_scoring = False
            args.physics_based = False
            args.mmff_minimization = False
            args.monte_carlo = False
            args.local_opt = False
            args.exhaustiveness = 1
            args.prepare_molecules = False
            args.population_size = 50  # Smaller population
        
        if args.monte_carlo:
            logger.info("\nRunning in Monte Carlo mode")
            args.algorithm = 'monte-carlo'
            args.exhaustiveness = 1

        if args.enhanced:
            logger.info("\nRunning with enhanced algorithms (slower but more accurate)")
            args.enhanced_scoring = True
            args.local_opt = True
            if args.population_size < 100:
                args.population_size = 100
        
        # Update status to reflect parameter choices
        update_status(
            output_dir,
            algorithm=args.algorithm,
            physics_based=args.physics_based,
            enhanced_scoring=args.enhanced_scoring,
            local_opt=args.local_opt,
            exhaustiveness=args.exhaustiveness,
            status="initializing"
        )
        
        # Prepare molecules if requested
        protein_path = args.protein
        ligand_path = args.ligand
        
        if args.prepare_molecules:
            logger.info("\nPreparing molecules for docking...")
            update_status(output_dir, status="preparing_molecules")
            
            # Prepare protein
            logger.info(f"Preparing protein {args.protein}...")
            prepared_protein = prepare_protein(
                args.protein, 
                output_file=temp_dir / f"prepared_{Path(args.protein).name}",
                ph=args.ph
            )
            
            # Prepare ligand
            logger.info(f"Preparing ligand {args.ligand}...")
            prepared_ligand = prepare_ligand(
                args.ligand,
                output_file=temp_dir / f"prepared_{Path(args.ligand).name}",
                n_conformers=5 if args.algorithm == 'genetic' else 1
            )
            
            # Update paths to prepared files
            protein_path = prepared_protein
            ligand_path = prepared_ligand
            
            # Update status
            update_status(
                output_dir, 
                prepared_protein=str(prepared_protein),
                prepared_ligand=str(prepared_ligand)
            )
        
        # Load protein
        logger.info(f"\nLoading protein from {protein_path}...")
        protein = Protein(protein_path)
        
        # Define active site
        if args.site:
            logger.info(f"Using provided active site center: {args.site}")
            protein.define_active_site(args.site, args.radius)
            update_status(
                output_dir,
                active_site_center=args.site,
                active_site_radius=args.radius
            )
        elif args.detect_pockets:
            logger.info("Detecting binding pockets...")
            update_status(output_dir, status="detecting_pockets")
            pockets = protein.detect_pockets()
            if pockets:
                logger.info(f"Found {len(pockets)} potential binding pockets")
                logger.info(f"Using largest pocket as active site")
                protein.define_active_site(pockets[0]['center'], pockets[0]['radius'])
                update_status(
                    output_dir,
                    active_site_center=pockets[0]['center'].tolist(),
                    active_site_radius=pockets[0]['radius'],
                    detected_pockets=len(pockets)
                )
            else:
                logger.info("No pockets detected, using whole protein")
                update_status(output_dir, no_pockets_detected=True)
        else:
            logger.info("No active site specified, using whole protein")
        
        # Check for flexible residues options
        if hasattr(args, 'auto_flex') and args.auto_flex:
            logger.info("Auto-flex option detected. Will attempt to automatically find flexible residues.")
            
        if args.auto_flex or args.flex_residues:
            logger.info("\nPreparing flexible protein configurations...")
            update_status(output_dir, status="preparing_flexible_residues")
            configs = prepare_protein_configs(protein, args, logger)
            
            if len(configs) > 1:
                logger.info(f"Using flexible protein configuration with {len(configs[1]['flexible_residues'])} flexible residues")
                protein = configs[1]['protein']  # Use the flexible configuration
                update_status(
                    output_dir,
                    flexible_residues=configs[1]['flexible_residues'],
                    flexible_residues_count=len(configs[1]['flexible_residues'])
                )
            else:
                logger.info("No flexible configuration available, using rigid protein")
                update_status(output_dir, flexible_residues_count=0)
        
        # Load ligand
        logger.info(f"\nLoading ligand from {ligand_path}...")
        ligand = Ligand(ligand_path)
        
        # Load reference ligand if provided
        reference_ligand = None
        if args.reference:
            logger.info(f"Loading reference ligand from {args.reference}...")
            reference_ligand = Ligand(args.reference)
            update_status(output_dir, reference_ligand=str(args.reference))
        
        # Setup hardware acceleration
        hybrid_manager = setup_hardware_acceleration(hw_config)
        
        # Apply MMFF minimization if requested
        if args.mmff_minimization and PHYSICS_AVAILABLE:
            logger.info("\nApplying MMFF94 force field minimization to ligand")
            update_status(output_dir, status="minimizing_ligand")
            minimizer = MMFFMinimization()
            ligand = minimizer.minimize_ligand(ligand)
            logger.info("Ligand minimization complete")
        elif args.mmff_minimization and not PHYSICS_AVAILABLE:
            logger.info("\nWarning: MMFF minimization requested but physics module not available. Skipping.")
        
        # Create scoring function based on hardware and requested type
        scoring_type = get_scoring_type_from_args(args)
        
        if scoring_type == 'physics' and PHYSICS_AVAILABLE:
            logger.info("\nUsing physics-based scoring function (MM-GBSA inspired)")
            scoring_function = PhysicsBasedScoring()
            update_status(output_dir, scoring_function="physics-based")
        else:
            # Use hardware-optimized scoring function
            scoring_function = create_optimized_scoring_function(hybrid_manager, scoring_type)
            
            if scoring_type == 'enhanced':
                logger.info("\nUsing enhanced scoring function with hardware acceleration")
                update_status(output_dir, scoring_function="enhanced")
            else:
                logger.info("\nUsing standard composite scoring function with hardware acceleration")
                update_status(output_dir, scoring_function="standard")
        
        # Initialize reporter
        reporter = DockingReporter(output_dir, args, timestamp=readable_date)
        
        # Get algorithm type and parameters
        algorithm_type = get_algorithm_type_from_args(args)
        algorithm_kwargs = get_algorithm_kwargs_from_args(args)
        
        # Pass output directory to algorithm kwargs for progress tracking
        algorithm_kwargs['output_dir'] = output_dir
        
        # Get advanced search if specified
        search_algorithm = None
        if args.advanced_search:
            # Use advanced search algorithm
            from .advanced_search import create_advanced_search_algorithm
            
            # Collect algorithm-specific parameters
            adv_search_kwargs = {'output_dir': output_dir}
            if args.advanced_search == 'gradient':
                adv_search_kwargs['gradient_step'] = args.gradient_step
                adv_search_kwargs['convergence_threshold'] = args.convergence_threshold
            elif args.advanced_search == 'replica-exchange':
                adv_search_kwargs['n_replicas'] = args.n_replicas
                adv_search_kwargs['temperatures'] = args.replica_temperatures
                adv_search_kwargs['exchange_steps'] = args.exchange_steps
            
            search_algorithm = create_advanced_search_algorithm(
                args.advanced_search,
                scoring_function,
                **adv_search_kwargs
            )
            logger.info(f"\nUsing advanced search algorithm: {args.advanced_search}")
        else:
            # Create the standard search algorithm
            search_algorithm = create_optimized_search_algorithm(
                hybrid_manager,
                algorithm_type,
                scoring_function,
                **algorithm_kwargs
            )
        
        # Update status before starting search
        update_status(
            output_dir, 
            status="searching",
            search_algorithm=algorithm_type,
            search_params=algorithm_kwargs
        )
        
        # Initialize results
        all_results = []
        
        # Make sure search algorithm has output directory
        if hasattr(search_algorithm, 'output_dir'):
            search_algorithm.output_dir = output_dir
        
        # Run the appropriate docking algorithm
        if args.reference and args.tethered_docking:
            logger.info(f"Using tethered reference-based docking with weight {args.tether_weight}...")
            all_results = search_algorithm.exact_reference_docking_with_tethering(
                protein,
                ligand,
                reference_ligand,
                tether_weight=args.tether_weight,
                skip_optimization=not args.local_opt # Skip if local_opt is FALSE
            )
        elif args.reference and args.exact_alignment:
            logger.info(f"Using exact alignment with reference structure...")
            all_results = search_algorithm.exact_reference_docking(
                protein,
                ligand,
                reference_ligand,
                skip_optimization=not args.local_opt # Skip if local_opt is FALSE
            )
        elif args.reference and not args.exact_alignment:
            logger.info(f"Using reference-guided docking...")
            all_results = search_algorithm.reference_guided_docking(
                protein,
                ligand,
                reference_ligand,
                skip_optimization=not args.local_opt # Skip if local_opt is FALSE
            )
        elif args.exhaustiveness > 1:
             # Ensemble docking needs separate handling if optimization is desired per run
             # For now, assume post-run optimization applies if --local-opt is set
             logger.info(f"\nRunning {args.exhaustiveness} independent docking runs...")
             all_results = hybrid_manager.run_ensemble_docking(
                 protein=protein,
                 ligand=ligand,
                 n_runs=args.exhaustiveness,
                 algorithm_type=algorithm_type,
                 **algorithm_kwargs
             )
             # Post-optimization logic below will handle these results if --local-opt is set
        elif algorithm_type == 'monte-carlo' and PHYSICS_AVAILABLE:
             # Monte Carlo usually includes optimization/sampling inherently
             logger.info(f"\nUsing Monte Carlo sampling with {args.mc_steps} steps at {args.temperature}K")
             mc_algorithm = MonteCarloSampling(
                 scoring_function,
                 temperature=args.temperature,
                 n_steps=args.mc_steps,
                 cooling_factor=args.cooling_factor,
                 output_dir=output_dir
             )
             all_results = mc_algorithm.run_sampling(protein, ligand)
             # Post-optimization logic might still refine MC results if --local-opt is set
        elif hasattr(args, 'enhanced') and args.enhanced:
             # Pass args to improve_rigid_docking so it knows about --local-opt
             logger.info("Using enhanced rigid docking algorithm...")
             all_results = search_algorithm.improve_rigid_docking(protein, ligand, args)
        else:
             logger.info("\nPerforming standard docking...")
             all_results = search_algorithm.search(protein, ligand)
             # Post-optimization logic below will handle these results if --local-opt is set
        
        # Update status after search completes
        update_status(
            output_dir, 
            status="post_processing",
            initial_poses_count=len(all_results)
        )
        
        # Apply local optimization to top poses if requested
        optimized_results = []  # Initialize variable
        if args.local_opt and all_results:
            logger.info("\nPerforming local optimization on top poses (enabled by --local-opt)...")
            update_status(output_dir, status="local_optimization")

            # Only optimize if we have results
            if all_results:
                # Optimize top 10 poses or fewer if we have less than 10 results
                poses_to_optimize = min(10, len(all_results))
                # Sort results before selecting top poses for optimization
                sorted_initial_results = sorted(all_results, key=lambda x: x[1])
                
                for i, (pose, score) in enumerate(sorted_initial_results[:poses_to_optimize]):
                    logger.info(f"Optimizing pose {i+1} (initial score: {score:.2f})...")
                    update_status(
                        output_dir,
                        optimizing_pose=i+1,
                        initial_score=score,
                        total_poses_to_optimize=poses_to_optimize
                    )

                    if args.mmff_minimization and PHYSICS_AVAILABLE:
                        # Use MMFF minimization in protein environment
                        logger.info(f"  Using MMFF minimization in protein environment")
                        minimizer = MMFFMinimization()
                        opt_pose = minimizer.minimize_pose(protein, pose)
                        opt_score = scoring_function.score(protein, opt_pose)
                        optimized_results.append((opt_pose, opt_score))
                        
                        # Save intermediate result
                        save_intermediate_result(opt_pose, opt_score, i, output_dir)
                    elif hasattr(search_algorithm, '_local_optimization'):
                        # Use built-in local optimization
                        opt_pose, opt_score = search_algorithm._local_optimization(pose, protein)
                        optimized_results.append((opt_pose, opt_score))
                        
                        # Save intermediate result
                        save_intermediate_result(opt_pose, opt_score, i, output_dir)
                    else:
                        # If no specific optimization method, keep the original
                        logger.info("  Warning: No specific local optimization method found for this setup. Keeping original pose.")
                        optimized_results.append((pose, score)) # Keep original if no method

                # Combine optimized results with the remaining unoptimized results
                # Ensure we don't add duplicates and maintain order based on original rank
                optimized_indices = {i for i in range(poses_to_optimize)}
                remaining_results = [r for i, r in enumerate(sorted_initial_results) if i not in optimized_indices]
                
                # Add optimized results first, then the rest
                all_results = optimized_results + remaining_results
                # Re-sort after optimization
                all_results.sort(key=lambda x: x[1])
                logger.info("Local optimization complete.")
                update_status(
                    output_dir,
                    status="optimization_complete",
                    optimized_poses_count=len(optimized_results)
                )
            else: # Should not happen if all_results check passed, but good practice
                 logger.info("  No results found to optimize.")

        # Sort final results if optimization happened or if it was skipped
        if all_results:
            all_results.sort(key=lambda x: x[1])
            update_status(
                output_dir,
                status="analyzing_results",
                final_poses_count=len(all_results),
                best_score=all_results[0][1]
            )

        # Apply analysis if requested
        if args.cluster_poses or args.analyze_interactions or args.classify_modes or \
            args.energy_decomposition or args.generate_analysis_report:
            
            # Only perform analysis if we have results
            if all_results:
                from .analysis import (PoseClusterer, InteractionFingerprinter, 
                                    BindingModeClassifier, EnergyDecomposition,
                                    DockingReportGenerator)
                
                logger.info("\nPerforming advanced analysis...")
                update_status(output_dir, status="analyzing")
                
                # Extract poses and scores
                poses = [pose for pose, _ in all_results]
                scores = [score for _, score in all_results]
                
                # Clustering
                clustering_results = None
                if args.cluster_poses:
                    logger.info("Clustering docking poses...")
                    update_status(output_dir, status="clustering")
                    clusterer = PoseClusterer(
                        method=args.clustering_method,
                        rmsd_cutoff=args.rmsd_cutoff
                    )
                    clustering_results = clusterer.cluster_poses(poses)
                    
                    # Print clustering summary
                    logger.info(f"Found {len(clustering_results['clusters'])} clusters")
                    for i, cluster in enumerate(clustering_results['clusters']):
                        logger.info(f"Cluster {i+1}: {len(cluster['members'])} poses, "
                            f"best score: {cluster['best_score']:.2f}")
                    
                    update_status(
                        output_dir,
                        clustering_complete=True,
                        clusters_count=len(clustering_results['clusters'])
                    )
                    
                    # Interaction analysis
                    if args.analyze_interactions:
                        logger.info("Analyzing protein-ligand interactions...")
                        update_status(output_dir, status="analyzing_interactions")
                        
                        # Create a file for interaction results
                        interaction_file = os.path.join(output_dir, "interaction_analysis.txt")
                        
                        with open(interaction_file, 'w') as f:
                            f.write("========================================\n")
                            f.write(" Protein-Ligand Interaction Analysis\n")
                            f.write("========================================\n\n")
                            
                            fingerprinter = InteractionFingerprinter(
                                interaction_types=args.interaction_types
                            )
                            
                            # Analyze top poses (up to 5)
                            poses_to_analyze = min(5, len(all_results))
                            
                            f.write(f"Analyzing top {poses_to_analyze} poses:\n\n")
                            
                            for i, (pose, score) in enumerate(all_results[:poses_to_analyze]):
                                f.write(f"Pose {i+1} (Score: {score:.2f})\n")
                                f.write("----------------------------\n")
                                
                                key_interactions = fingerprinter.analyze_key_interactions(protein, pose)
                                for interaction in key_interactions:
                                    f.write(f"  {interaction}\n")
                                
                                f.write("\n")  # Add space between poses
                                
                                # Also log to console
                                logger.info(f"\nInteractions for pose {i+1} (score: {score:.2f}):")
                                for interaction in key_interactions:
                                    logger.info(f"  {interaction}")
                        
                        logger.info(f"\nInteraction analysis saved to: {interaction_file}")
        
                        #if args.analyze_interactions:
                         #   logger.info("Analyzing protein-ligand interactions...")
                         #   update_status(output_dir, status="analyzing_interactions")
                          #  fingerprinter = InteractionFingerprinter(
                          #      interaction_types=args.interaction_types
                          #  )
                            # Analyze top poses (up to 5)
                          #  poses_to_analyze = min(5, len(all_results))
                          #  for i, (pose, score) in enumerate(all_results[:poses_to_analyze]):
                          #      logger.info(f"\nInteractions for pose {i+1} (score: {score:.2f}):")
                          #      key_interactions = fingerprinter.analyze_key_interactions(protein, pose)
                          #      for interaction in key_interactions:
                          #          logger.info(f"  {interaction}")
                    
                    # Binding mode classification
                    if args.classify_modes or args.discover_modes:
                        logger.info("Analyzing binding modes...")
                        update_status(output_dir, status="analyzing_binding_modes")
                        classifier = BindingModeClassifier()
                        
                        if args.discover_modes:
                            discovered_modes = classifier.discover_modes(
                                protein, poses, n_modes=args.n_modes
                            )
                            logger.info(f"Discovered {len(discovered_modes)} binding modes")
                            for i, mode in enumerate(discovered_modes):
                                logger.info(f"Mode {i+1}: {mode['count']} poses, "
                                    f"best score: {mode['best_score']:.2f}")
                                
                        if args.classify_modes:
                            # Analyze top 10 poses or fewer if we have less
                            poses_to_classify = min(10, len(all_results))
                            for i, (pose, score) in enumerate(all_results[:poses_to_classify]):
                                mode = classifier.classify_pose(protein, pose)
                                logger.info(f"Pose {i+1} (score: {score:.2f}): {mode}")
                    
                    # Energy decomposition
                    energy_decomposition = None
                    if args.energy_decomposition and all_results:
                        logger.info("Performing energy decomposition analysis...")
                        update_status(output_dir, status="energy_decomposition")
                        decomposer = EnergyDecomposition(scoring_function)
                        
                        # Analyze top pose
                        top_pose = all_results[0][0]
                        energy_decomposition = decomposer.decompose_energy(protein, top_pose)
                        
                        logger.info("\nEnergy components for top pose:")
                        for component, value in energy_decomposition.items():
                            logger.info(f"  {component}: {value:.2f}")
                            
                        if args.per_residue_energy:
                            logger.info("\nTop residue contributions:")
                            res_contributions = decomposer.residue_contributions(protein, top_pose)
                            for res, value in res_contributions[:5]:
                                logger.info(f"  {res}: {value:.2f}")
                
                # Generate report
                if args.generate_analysis_report:
                    logger.info("Generating comprehensive docking report...")
                    update_status(output_dir, status="generating_report")
                    report_generator = DockingReportGenerator(
                        report_format=args.analysis_report_format,
                        include_sections=args.analysis_report_sections
                    )
                    
                    report_file = os.path.join(output_dir, f"docking_report.{args.analysis_report_format}")
                    report_generator.generate_report(
                        protein, poses, scores, report_file,
                        clustering_results=clustering_results,
                        energy_decomposition=energy_decomposition
                    )
                    logger.info(f"Report generated: {report_file}")
            else:
                logger.info("\nSkipping analysis as no valid docking solutions were found.")

        # Extract energy components for reporting if possible
        try:
            if all_results:
                logger.info("Extracting energy components for detailed reporting...")
                update_status(output_dir, status="extracting_energy_components")
                energy_breakdown = reporter.extract_energy_components(
                    scoring_function, 
                    protein, 
                    [pose for pose, _ in all_results[:min(20, len(all_results))]]
                )
                reporter.add_results(all_results, energy_breakdown)
            else:
                reporter.add_results([])  # Add empty results
        except Exception as e:
            logger.info(f"Warning: Could not extract energy components: {e}")
            reporter.add_results(all_results)
        
        # Take the top results (avoid duplicates)
        unique_results = []
        seen_scores = set()
        
        for pose, score in all_results:
            # Round score to avoid floating point comparison issues
            rounded_score = round(score, 4)
            if rounded_score not in seen_scores:
                unique_results.append((pose, score))
                seen_scores.add(rounded_score)
                
                # Keep at most 20 poses
                if len(unique_results) >= 20:
                    break
        
        # Calculate elapsed time here, before potential early returns
        elapsed_time = time.time() - start_time
        
        # Check if we have any results before saving
        if not unique_results:
            logger.info("\nWarning: No valid docking solutions found.")
            update_status(output_dir, status="completed", success=False)
        else:
            # Save results if we have any
            logger.info(f"\nDocking completed successfully!")
            logger.info(f"Saving results to {output_dir}...")
            update_status(output_dir, status="saving_results")
            
            # Pass flexible residues to save_docking_results if they exist
            flexible_residues = None
            if hasattr(protein, 'flexible_residues') and protein.flexible_residues:
                flexible_residues = protein.flexible_residues
                logger.info(f"Found {len(flexible_residues)} flexible residues to include in output")
            else:
                logger.info("No flexible residues found on protein object")
            
            # Save docking results
            save_docking_results(unique_results, output_dir, flexible_residues=flexible_residues)
            
            # Validate against reference if provided
            if hasattr(args, 'reference') and args.reference and not args.exact_alignment:
                logger.info("Validating against reference structure...")
                update_status(output_dir, status="validating")
                validation_results = validate_against_reference(args, unique_results, output_dir)
                # Add validation results to reporter
                reporter.add_validation_results(validation_results)
        
        # Write detailed results to text file
        write_results_to_txt(
            results=unique_results,
            output_dir=output_dir,
            elapsed_time=elapsed_time,
            protein_path=args.protein,
            ligand_path=args.ligand,
            algorithm=algorithm_type,
            iterations=args.iterations if algorithm_type != 'monte-carlo' else args.mc_steps,
            logger=logger
        )
        
        # Print summary
        logger.info(f"\nDocking completed in {elapsed_time:.1f} seconds")
        if unique_results:
            logger.info(f"Best score: {unique_results[0][1]:.2f}")
            logger.info(f"Results saved to: {output_dir}")
        else:
            logger.info("No valid docking solutions found.")
        logger.info(f"============================================================")
        
        # Generate comprehensive reports
        logger.info("Generating comprehensive docking reports...")
        update_status(output_dir, status="generating_reports")
        if hasattr(args, 'report_format') and args.report_format != 'all':
            # Generate only the requested format
            if args.report_format == 'text':
                reporter.generate_detailed_report()
            elif args.report_format == 'csv':
                reporter.generate_csv_report()
            elif args.report_format == 'json':
                reporter.generate_json_report()
            elif args.report_format == 'html':
                reporter.generate_html_report()
        else:
            # Generate all report formats
            reporter.generate_detailed_report()
            reporter.generate_csv_report()
            reporter.generate_json_report()
            
            # Generate HTML report with visualizations if plots are not skipped
            if not (hasattr(args, 'skip_plots') and args.skip_plots):
                html_report = reporter.generate_html_report()
                logger.info(f"Comprehensive HTML report with visualizations saved to: {html_report}")
        
        # Final status update
        update_status(
            output_dir,
            status="completed",
            success=True,
            elapsed_time=elapsed_time,
            best_score=unique_results[0][1] if unique_results else None,
            final_poses_count=len(unique_results)
        )
                
    except Exception as e:
        # Calculate elapsed time if there was an error
        elapsed_time = time.time() - start_time
        
        # Try to get logger, if available
        try:
            logger.error(f"\nError during docking: {str(e)}")
            logger.error(traceback.format_exc())
        except:
            print(f"\nError during docking: {str(e)}")
            traceback.print_exc()
            
        # Try to update status file if output_dir exists
        if output_dir:
            try:
                update_status(
                    output_dir,
                    status="error",
                    error_message=str(e),
                    traceback=traceback.format_exc(),
                    elapsed_time=elapsed_time
                )
            except:
                pass
                
        print(f"\nDocking failed after {elapsed_time:.1f} seconds")
        return_code = 1
    
    finally:
        # Clean up temporary files
        if temp_dir is not None and hasattr(args, 'prepare_molecules') and args.prepare_molecules:
            import shutil
            try:
                shutil.rmtree(temp_dir, ignore_errors=True)
                if 'logger' in locals():
                    logger.info(f"Cleaned up temporary directory: {temp_dir}")
                else:
                    print(f"Cleaned up temporary directory: {temp_dir}")
            except Exception as e:
                if 'logger' in locals():
                    logger.warning(f"Failed to clean up temporary directory: {e}")
                else:
                    print(f"Failed to clean up temporary directory: {e}")
        
        # Clean up hardware resources
        if hybrid_manager is not None:
            try:
                hybrid_manager.cleanup()
                if 'logger' in locals():
                    logger.info("Hardware resources cleaned up")
                else:
                    print("Hardware resources cleaned up")
            except Exception as e:
                if 'logger' in locals():
                    logger.warning(f"Failed to clean up hardware resources: {e}")
                else:
                    print(f"Failed to clean up hardware resources: {e}")
            
        # Final log message
        if 'logger' in locals():
            logger.info(f"============================================================")
        else:
            print(f"============================================================")
            
        return return_code

if __name__ == "__main__":
    main()
