"""
Ring numbering system for 16-membered macrolactones.
"""
from rdkit import Chem
from typing import Dict, Tuple, Optional, List


def find_lactone_carbon(mol: Chem.Mol) -> Optional[int]:
    """
    Find the lactone carbonyl carbon in the 16-membered ring.
    Uses SMARTS pattern: [r16;#8][r16;#6](=[#8]) to ensure both O and C are on the ring.
    
    Args:
        mol: RDKit molecule object
    
    Returns:
        Atom index of the lactone carbonyl carbon (C=O in the ring)
        Returns None if not found
    """
    # First, get the 16-membered ring atoms
    ring_atoms = get_ring_atoms(mol)
    if len(ring_atoms) == 0:
        return None
    
    ring_atoms_set = set(ring_atoms)
    
    # SMARTS pattern to match ester group where both O and C are on the ring
    # [r16;#8] = ring atom that is oxygen
    # [r16;#6](=[#8]) = ring atom that is carbon with double bond to oxygen
    lactone_pattern = Chem.MolFromSmarts("[r16;#8][r16;#6](=[#8])")
    
    if lactone_pattern is None:
        return None
    
    matches = mol.GetSubstructMatches(lactone_pattern)
    
    if not matches:
        return None
    
    # Get the first match
    # matches[0] contains: (ring_O, ring_C, O_double)
    # We want the carbon (index 1) which should be in the ring
    for match in matches:
        if len(match) >= 2:
            carbonyl_carbon = match[1]  # The carbonyl carbon
            # Double check that it's in the ring (should be guaranteed by SMARTS, but verify)
            if carbonyl_carbon in ring_atoms_set:
                return carbonyl_carbon
    
    return None


def get_ring_atoms(mol: Chem.Mol) -> List[int]:
    """
    Get all atoms in the 16-membered ring.
    
    Args:
        mol: RDKit molecule object
    
    Returns:
        List of atom indices in the 16-membered ring
    """
    # Find atoms in 16-membered rings
    # Note: Use r16 directly, not r{16} which is not supported in all RDKit versions
    ring_pattern = Chem.MolFromSmarts("[r16]")
    
    if ring_pattern is None:
        return []
    
    matches = mol.GetSubstructMatches(ring_pattern)
    
    # Extract unique atom indices
    ring_atoms = list(set([match[0] for match in matches]))
    
    return ring_atoms


def order_ring_atoms_clockwise(mol: Chem.Mol, ring_atoms: List[int], start_atom: int) -> List[int]:
    """
    Order ring atoms in a clockwise manner starting from a specific atom.
    
    Args:
        mol: RDKit molecule object
        ring_atoms: List of atom indices in the ring
        start_atom: Starting atom index
    
    Returns:
        Ordered list of atom indices
    """
    if start_atom not in ring_atoms:
        raise ValueError("Start atom is not in the ring")
    
    ordered = [start_atom]
    visited = {start_atom}
    current = start_atom
    
    while len(ordered) < len(ring_atoms):
        atom = mol.GetAtomWithIdx(current)
        
        # Find the next ring atom among neighbors
        next_atom = None
        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            if neighbor_idx in ring_atoms and neighbor_idx not in visited:
                next_atom = neighbor_idx
                break
        
        if next_atom is None:
            # If we can't find a forward path, we might need to go backwards from start
            if len(ordered) < len(ring_atoms):
                # Try going backwards from the start
                atom = mol.GetAtomWithIdx(start_atom)
                for neighbor in atom.GetNeighbors():
                    neighbor_idx = neighbor.GetIdx()
                    if neighbor_idx in ring_atoms and neighbor_idx not in visited:
                        # Insert at the beginning
                        ordered.insert(0, neighbor_idx)
                        visited.add(neighbor_idx)
                        current = neighbor_idx
                        break
                else:
                    break
            else:
                break
        else:
            ordered.append(next_atom)
            visited.add(next_atom)
            current = next_atom
    
    return ordered


def assign_ring_numbering(mol: Chem.Mol) -> Dict[int, int]:
    """
    Assign fixed numbering (1-16) to atoms in the 16-membered ring.
    Numbering starts from the lactone carbonyl carbon (position 1) and proceeds clockwise.
    
    Args:
        mol: RDKit molecule object
    
    Returns:
        Dictionary mapping atom index to ring position (1-16)
        Returns empty dict if no valid ring found
    """
    # Find the lactone carbon (will be position 1)
    lactone_c = find_lactone_carbon(mol)
    
    if lactone_c is None:
        print("Warning: Could not find lactone carbon")
        return {}
    
    # Get all 16-ring atoms
    ring_atoms = get_ring_atoms(mol)
    
    if len(ring_atoms) != 16:
        print(f"Warning: Expected 16 ring atoms, found {len(ring_atoms)}")
        # Still proceed if we found some ring atoms
        if len(ring_atoms) == 0:
            return {}
    
    # Order atoms clockwise starting from lactone carbon
    ordered_atoms = order_ring_atoms_clockwise(mol, ring_atoms, lactone_c)
    
    # Create numbering dictionary
    numbering = {}
    for position, atom_idx in enumerate(ordered_atoms, start=1):
        numbering[atom_idx] = position
    
    return numbering


def get_ring_position(numbering: Dict[int, int], atom_idx: int) -> Optional[int]:
    """
    Get the ring position for a given atom index.
    
    Args:
        numbering: Dictionary from assign_ring_numbering
        atom_idx: Atom index
    
    Returns:
        Ring position (1-16) or None if not in ring
    """
    return numbering.get(atom_idx)


def validate_numbering(mol: Chem.Mol, numbering: Dict[int, int]) -> bool:
    """
    Validate that the numbering is correct.
    
    Args:
        mol: RDKit molecule object
        numbering: Dictionary from assign_ring_numbering
    
    Returns:
        True if valid, False otherwise
    """
    if len(numbering) != 16:
        return False
    
    # Check that positions are 1-16
    positions = set(numbering.values())
    if positions != set(range(1, 17)):
        return False
    
    # Check that position 1 is the lactone carbon
    lactone_c = find_lactone_carbon(mol)
    if lactone_c is None or numbering.get(lactone_c) != 1:
        return False
    
    return True

