"""
Docking Utilities for OdinFold

Utility functions for preparing proteins and ligands for molecular docking,
binding site detection, and result analysis.
"""

import numpy as np
import torch
from typing import Dict, List, Optional, Tuple, Union
from pathlib import Path
import logging
import tempfile

try:
    from Bio.PDB import PDBParser, PDBIO, Select
    from Bio.PDB.Structure import Structure
    BIOPYTHON_AVAILABLE = True
except ImportError:
    BIOPYTHON_AVAILABLE = False
    logging.warning("BioPython not available. Some functionality will be limited.")

logger = logging.getLogger(__name__)


def prepare_protein_for_docking(coordinates: torch.Tensor, 
                              sequence: str,
                              chain_id: str = "A",
                              output_path: Optional[Union[str, Path]] = None) -> str:
    """
    Prepare protein structure for docking from OdinFold coordinates.
    
    Args:
        coordinates: Protein coordinates [seq_len, 3]
        sequence: Amino acid sequence
        chain_id: Chain identifier
        output_path: Optional output PDB file path
        
    Returns:
        PDB content as string or path to saved file
    """
    
    if len(coordinates) != len(sequence):
        raise ValueError("Coordinates and sequence length mismatch")
    
    # Convert to numpy for easier manipulation
    coords = coordinates.detach().cpu().numpy()
    
    # Generate PDB content
    pdb_lines = []
    
    # Header
    pdb_lines.append("HEADER    PROTEIN STRUCTURE FROM ODINFOLD")
    pdb_lines.append("REMARK   Generated by OdinFold for molecular docking")
    
    # Atoms (CA only for simplicity)
    for i, (aa, coord) in enumerate(zip(sequence, coords)):
        atom_line = (
            f"ATOM  {i+1:5d}  CA  {aa_three_letter(aa)} {chain_id}{i+1:4d}    "
            f"{coord[0]:8.3f}{coord[1]:8.3f}{coord[2]:8.3f}  1.00 20.00           C  "
        )
        pdb_lines.append(atom_line)
    
    pdb_lines.append("END")
    
    pdb_content = "\n".join(pdb_lines)
    
    # Save to file if path provided
    if output_path:
        output_path = Path(output_path)
        with open(output_path, 'w') as f:
            f.write(pdb_content)
        return str(output_path)
    
    return pdb_content


def prepare_ligand_for_docking(ligand_data: Dict[str, torch.Tensor],
                             output_path: Optional[Union[str, Path]] = None) -> str:
    """
    Prepare ligand for docking from OdinFold ligand data.
    
    Args:
        ligand_data: Ligand data dictionary with positions and features
        output_path: Optional output file path
        
    Returns:
        File path or content string
    """
    
    positions = ligand_data['positions'].detach().cpu().numpy()
    atom_features = ligand_data['atom_features'].detach().cpu().numpy()
    
    # Generate simple MOL2 format
    mol2_lines = []
    
    # Header
    mol2_lines.append("@<TRIPOS>MOLECULE")
    mol2_lines.append("ligand")
    mol2_lines.append(f"{len(positions)} 0 0 0 0")
    mol2_lines.append("SMALL")
    mol2_lines.append("GASTEIGER")
    mol2_lines.append("")
    
    # Atoms
    mol2_lines.append("@<TRIPOS>ATOM")
    for i, (pos, features) in enumerate(zip(positions, atom_features)):
        atom_type = int(features[0])  # Atomic number
        element = atomic_number_to_symbol(atom_type)
        
        atom_line = (
            f"{i+1:7d} {element}{i+1:<4d} {pos[0]:10.4f} {pos[1]:10.4f} {pos[2]:10.4f} "
            f"{element:<5s} 1    LIG    0.0000"
        )
        mol2_lines.append(atom_line)
    
    mol2_content = "\n".join(mol2_lines)
    
    # Save to file if path provided
    if output_path:
        output_path = Path(output_path)
        with open(output_path, 'w') as f:
            f.write(mol2_content)
        return str(output_path)
    
    return mol2_content


def detect_binding_site(protein_coords: torch.Tensor,
                       ligand_coords: Optional[torch.Tensor] = None,
                       method: str = "geometric_center") -> Dict[str, float]:
    """
    Detect binding site for molecular docking.
    
    Args:
        protein_coords: Protein coordinates [seq_len, 3]
        ligand_coords: Optional ligand coordinates for reference [num_atoms, 3]
        method: Detection method ("geometric_center", "cavity", "ligand_based")
        
    Returns:
        Binding site parameters (center and size)
    """
    
    coords = protein_coords.detach().cpu().numpy()
    
    if method == "geometric_center":
        # Use protein geometric center
        center = coords.mean(axis=0)
        
        # Calculate reasonable box size based on protein dimensions
        protein_span = coords.max(axis=0) - coords.min(axis=0)
        box_size = np.minimum(protein_span * 0.5, 25.0)  # Max 25Å
        box_size = np.maximum(box_size, 15.0)  # Min 15Å
        
    elif method == "ligand_based" and ligand_coords is not None:
        # Use ligand position as reference
        ligand_coords_np = ligand_coords.detach().cpu().numpy()
        center = ligand_coords_np.mean(axis=0)
        
        # Box size based on ligand span plus buffer
        ligand_span = ligand_coords_np.max(axis=0) - ligand_coords_np.min(axis=0)
        box_size = ligand_span + 10.0  # 10Å buffer
        box_size = np.maximum(box_size, 15.0)  # Min 15Å
        
    elif method == "cavity":
        # Simple cavity detection (find largest empty space)
        center, box_size = detect_largest_cavity(coords)
        
    else:
        raise ValueError(f"Unknown binding site detection method: {method}")
    
    return {
        'center_x': float(center[0]),
        'center_y': float(center[1]),
        'center_z': float(center[2]),
        'size_x': float(box_size[0]),
        'size_y': float(box_size[1]),
        'size_z': float(box_size[2])
    }


def detect_largest_cavity(coords: np.ndarray, grid_spacing: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
    """
    Detect largest cavity in protein structure using grid-based approach.
    
    Args:
        coords: Protein coordinates [seq_len, 3]
        grid_spacing: Grid spacing for cavity detection
        
    Returns:
        Cavity center and size
    """
    
    # Define grid bounds
    min_coords = coords.min(axis=0) - 5.0
    max_coords = coords.max(axis=0) + 5.0
    
    # Create grid
    x_grid = np.arange(min_coords[0], max_coords[0], grid_spacing)
    y_grid = np.arange(min_coords[1], max_coords[1], grid_spacing)
    z_grid = np.arange(min_coords[2], max_coords[2], grid_spacing)
    
    # Find grid points far from protein atoms
    cavity_points = []
    min_distance_threshold = 3.0  # Minimum distance from protein
    
    for x in x_grid[::2]:  # Subsample for speed
        for y in y_grid[::2]:
            for z in z_grid[::2]:
                point = np.array([x, y, z])
                
                # Calculate minimum distance to any protein atom
                distances = np.linalg.norm(coords - point, axis=1)
                min_distance = distances.min()
                
                if min_distance > min_distance_threshold:
                    cavity_points.append(point)
    
    if not cavity_points:
        # Fallback to geometric center
        center = coords.mean(axis=0)
        size = np.array([20.0, 20.0, 20.0])
        return center, size
    
    cavity_points = np.array(cavity_points)
    
    # Find largest connected cavity (simplified)
    center = cavity_points.mean(axis=0)
    
    # Calculate cavity size
    cavity_span = cavity_points.max(axis=0) - cavity_points.min(axis=0)
    size = np.minimum(cavity_span, 25.0)  # Max 25Å
    size = np.maximum(size, 15.0)  # Min 15Å
    
    return center, size


def aa_three_letter(aa_one: str) -> str:
    """Convert single-letter amino acid code to three-letter code."""
    
    aa_map = {
        'A': 'ALA', 'R': 'ARG', 'N': 'ASN', 'D': 'ASP', 'C': 'CYS',
        'Q': 'GLN', 'E': 'GLU', 'G': 'GLY', 'H': 'HIS', 'I': 'ILE',
        'L': 'LEU', 'K': 'LYS', 'M': 'MET', 'F': 'PHE', 'P': 'PRO',
        'S': 'SER', 'T': 'THR', 'W': 'TRP', 'Y': 'TYR', 'V': 'VAL'
    }
    
    return aa_map.get(aa_one.upper(), 'UNK')


def atomic_number_to_symbol(atomic_number: int) -> str:
    """Convert atomic number to element symbol."""
    
    element_map = {
        1: 'H', 2: 'He', 3: 'Li', 4: 'Be', 5: 'B', 6: 'C', 7: 'N', 8: 'O',
        9: 'F', 10: 'Ne', 11: 'Na', 12: 'Mg', 13: 'Al', 14: 'Si', 15: 'P',
        16: 'S', 17: 'Cl', 18: 'Ar', 19: 'K', 20: 'Ca'
    }
    
    return element_map.get(atomic_number, 'X')


def calculate_rmsd(coords1: torch.Tensor, coords2: torch.Tensor) -> float:
    """
    Calculate RMSD between two coordinate sets.
    
    Args:
        coords1: First coordinate set [N, 3]
        coords2: Second coordinate set [N, 3]
        
    Returns:
        RMSD value
    """
    
    if coords1.shape != coords2.shape:
        raise ValueError("Coordinate shapes must match")
    
    diff = coords1 - coords2
    squared_diff = (diff ** 2).sum(dim=-1)
    rmsd = torch.sqrt(squared_diff.mean())
    
    return float(rmsd)


def analyze_binding_pocket(protein_coords: torch.Tensor,
                          ligand_coords: torch.Tensor,
                          cutoff: float = 5.0) -> Dict:
    """
    Analyze binding pocket properties.
    
    Args:
        protein_coords: Protein coordinates [seq_len, 3]
        ligand_coords: Ligand coordinates [num_atoms, 3]
        cutoff: Distance cutoff for pocket definition
        
    Returns:
        Binding pocket analysis
    """
    
    # Convert to numpy
    protein_np = protein_coords.detach().cpu().numpy()
    ligand_np = ligand_coords.detach().cpu().numpy()
    
    # Calculate distances
    distances = np.linalg.norm(
        protein_np[:, np.newaxis, :] - ligand_np[np.newaxis, :, :], 
        axis=2
    )
    
    # Find pocket residues
    min_distances = distances.min(axis=1)
    pocket_residues = min_distances <= cutoff
    pocket_indices = np.where(pocket_residues)[0]
    
    # Calculate pocket properties
    if len(pocket_indices) > 0:
        pocket_coords = protein_np[pocket_indices]
        pocket_center = pocket_coords.mean(axis=0)
        pocket_span = pocket_coords.max(axis=0) - pocket_coords.min(axis=0)
        pocket_volume = np.prod(pocket_span)  # Approximate volume
    else:
        pocket_center = np.array([0.0, 0.0, 0.0])
        pocket_span = np.array([0.0, 0.0, 0.0])
        pocket_volume = 0.0
    
    # Calculate ligand properties
    ligand_center = ligand_np.mean(axis=0)
    ligand_span = ligand_np.max(axis=0) - ligand_np.min(axis=0)
    
    # Binding complementarity (how well ligand fits pocket)
    center_distance = np.linalg.norm(pocket_center - ligand_center)
    size_ratio = np.prod(ligand_span) / max(pocket_volume, 1e-6)
    
    return {
        'pocket_residues': pocket_indices.tolist(),
        'pocket_size': len(pocket_indices),
        'pocket_center': pocket_center.tolist(),
        'pocket_span': pocket_span.tolist(),
        'pocket_volume': float(pocket_volume),
        'ligand_center': ligand_center.tolist(),
        'ligand_span': ligand_span.tolist(),
        'center_distance': float(center_distance),
        'size_ratio': float(size_ratio),
        'complementarity_score': float(1.0 / (1.0 + center_distance + abs(1.0 - size_ratio)))
    }
