app.py

import gradio as gr
import re
import pandas as pd
from io import StringIO
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem, Draw
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from io import BytesIO

def is_peptide(smiles):
    """Check if the SMILES represents a peptide by looking for peptide bonds"""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return False
        
    # Look for peptide bonds: NC(=O) pattern
    peptide_bond_pattern = Chem.MolFromSmarts('[NH][C](=O)')
    if mol.HasSubstructMatch(peptide_bond_pattern):
        return True
        
    # Look for N-methylated peptide bonds: N(C)C(=O) pattern
    n_methyl_pattern = Chem.MolFromSmarts('[N;H0;$(NC)](C)[C](=O)')
    if mol.HasSubstructMatch(n_methyl_pattern):
        return True
        
    # Look for ester bonds in cyclic depsipeptides: OC(=O) pattern
    ester_bond_pattern = Chem.MolFromSmarts('O[C](=O)')
    if mol.HasSubstructMatch(ester_bond_pattern):
        return True
        
    return False

def remove_nested_branches(smiles):
    """Remove nested branches from SMILES string"""
    result = ''
    depth = 0
    for char in smiles:
        if char == '(':
            depth += 1
        elif char == ')':
            depth -= 1
        elif depth == 0:
            result += char
    return result

def identify_linkage_type(segment):
    """
    Identify the type of linkage between residues
    Returns: tuple (type, is_n_methylated)
    """
    if 'OC(=O)' in segment:
        return ('ester', False)
    elif 'N(C)C(=O)' in segment:
        return ('peptide', True)  # N-methylated peptide bond
    elif 'NC(=O)' in segment:
        return ('peptide', False)  # Regular peptide bond
    return (None, False)
def identify_residue(segment, next_segment=None, prev_segment=None):
    """
    Identify amino acid residues with modifications and special handling for Proline
    Returns: tuple (residue, modifications)
    """
    modifications = []
    
    # Check for modifications in the next segment
    if next_segment:
        if 'N(C)C(=O)' in next_segment:
            modifications.append('N-Me')
        if 'OC(=O)' in next_segment:
            modifications.append('O-linked')

    # Special case for Proline - check for CCCN pattern and its cyclization
    # Proline can appear in several patterns due to its cyclic nature
    if any(pattern in segment for pattern in ['CCCN2', 'N2CCC', '[C@@H]2CCCN2', 'CCCN1', 'N1CCC']):
        return ('Pro', modifications)
    
    # Check if this segment is part of a Proline ring by looking at context
    if prev_segment and next_segment:
        if ('CCC' in segment and 'N' in next_segment) or ('N' in segment and 'CCC' in prev_segment):
            combined = prev_segment + segment + next_segment
            if re.search(r'CCCN.*C\(=O\)', combined):
                return ('Pro', modifications)

    # Aromatic amino acids
    if 'Cc2ccccc2' in segment or 'c1ccccc1' in segment:  
        return ('Phe', modifications)
    if 'c2ccc(O)cc2' in segment:  
        return ('Tyr', modifications)
    if 'c1c[nH]c2ccccc12' in segment:  
        return ('Trp', modifications)
    if 'c1cnc[nH]1' in segment:  
        return ('His', modifications)
        
    # Branched chain amino acids
    if 'CC(C)C[C@H]' in segment or 'CC(C)C[C@@H]' in segment:  
        return ('Leu', modifications)
    if '[C@H](CC(C)C)' in segment or '[C@@H](CC(C)C)' in segment:  
        return ('Leu', modifications)
    if 'C(C)C' in segment and not any(pat in segment for pat in ['CC(C)C', 'C(C)C[C@H]', 'C(C)C[C@@H]']):
        return ('Val', modifications)
    if 'C(C)C[C@H]' in segment or 'C(C)C[C@@H]' in segment:  
        return ('Ile', modifications)
        
    # Small/polar amino acids
    if ('[C@H](C)' in segment or '[C@@H](C)' in segment) and 'C(C)C' not in segment:
        return ('Ala', modifications)
    if '[C@H](CO)' in segment:
        return ('Ser', modifications)
    if '[C@H](C(C)O)' in segment or '[C@@H](C(C)O)' in segment:
        return ('Thr', modifications)
    if '[C@H]' in segment and not any(pat in segment for pat in ['C(C)', 'CC', 'O', 'N', 'S']):
        return ('Gly', modifications)
        
    # Rest of amino acids remain the same...
    # [Previous code for other amino acids]
    
    return (None, modifications)
def parse_peptide(smiles):
    """
    Parse peptide sequence with enhanced Proline recognition
    """
    # Split on peptide bonds while preserving cycle numbers
    bond_pattern = r'(NC\(=O\)|N\(C\)C\(=O\)|N\dC\(=O\)|OC\(=O\))'
    segments = re.split(bond_pattern, smiles)
    segments = [s for s in segments if s]
    
    sequence = []
    i = 0
    while i < len(segments):
        segment = segments[i]
        next_segment = segments[i+1] if i+1 < len(segments) else None
        prev_segment = segments[i-1] if i > 0 else None
        
        # Skip pure bond patterns
        if re.match(r'.*C\(=O\)$', segment):
            i += 1
            continue
            
        residue, modifications = identify_residue(segment, next_segment, prev_segment)
        if residue:
            # Format residue with modifications
            formatted_residue = residue
            if modifications:
                formatted_residue += f"({','.join(modifications)})"
            sequence.append(formatted_residue)
        
        i += 1
    
    is_cyclic = is_cyclic_peptide(smiles)
    
    # Print debug information
    print("\nDetailed Analysis:")
    print("Segments:", segments)
    print("Found sequence:", sequence)
    
    # Format the final sequence
    if is_cyclic:
        return f"cyclo({'-'.join(sequence)})"
    return '-'.join(sequence)

def is_cyclic_peptide(smiles):
    """
    Determine if SMILES represents a cyclic peptide by checking:
    1. Proper cycle number pairing
    2. Presence of peptide bonds between cycle points
    3. Distinguishing between aromatic rings and peptide cycles
    """
    cycle_info = {}
    
    # Find all cycle numbers and their contexts
    for match in re.finditer(r'(\w{3})?(\d)(\w{3})?', smiles):
        number = match.group(2)
        pre_context = match.group(1) or ''
        post_context = match.group(3) or ''
        position = match.start(2)
        
        if number not in cycle_info:
            cycle_info[number] = []
        cycle_info[number].append({
            'position': position,
            'pre_context': pre_context,
            'post_context': post_context,
            'full_context': smiles[max(0, position-3):min(len(smiles), position+4)]
        })
    
    # Check each cycle
    peptide_cycles = []
    aromatic_cycles = []
    
    for number, occurrences in cycle_info.items():
        if len(occurrences) != 2:  # Must have exactly 2 occurrences
            continue
            
        start, end = occurrences[0]['position'], occurrences[1]['position']
        
        # Get the segment between cycle points
        segment = smiles[start:end+1]
        clean_segment = remove_nested_branches(segment)
        
        # Check if this is an aromatic ring
        is_aromatic = any(context['full_context'].count('c') >= 2 for context in occurrences)
        
        # Check if this is a peptide cycle
        has_peptide_bond = 'NC(=O)' in segment or 'N2C(=O)' in segment
        
        if is_aromatic:
            aromatic_cycles.append(number)
        elif has_peptide_bond:
            peptide_cycles.append(number)
    
    return len(peptide_cycles) > 0, peptide_cycles, aromatic_cycles

def analyze_single_smiles(smiles):
    """Analyze a single SMILES string"""
    try:
        is_cyclic, peptide_cycles, aromatic_cycles = is_cyclic_peptide(smiles)
        sequence = parse_peptide(smiles)
        
        details = {
            #'SMILES': smiles,
            'Sequence': sequence,
            'Is Cyclic': 'Yes' if is_cyclic else 'No',
            #'Peptide Cycles': ', '.join(peptide_cycles) if peptide_cycles else 'None',
            #'Aromatic Cycles': ', '.join(aromatic_cycles) if aromatic_cycles else 'None'
        }
        return details
        
    except Exception as e:
        return {
            #'SMILES': smiles,
            'Sequence': f'Error: {str(e)}',
            'Is Cyclic': 'Error',
            #'Peptide Cycles': 'Error',
            #'Aromatic Cycles': 'Error'
        }

def annotate_cyclic_structure(mol, sequence):
    """Create annotated 2D structure with clear, non-overlapping residue labels"""
    # Generate 2D coordinates
    AllChem.Compute2DCoords(mol)
    
    # Create drawer with larger size for annotations
    drawer = Draw.rdMolDraw2D.MolDraw2DCairo(2000, 2000)  # Even larger size
    
    # Get residue list
    if sequence.startswith('cyclo('):
        residues = sequence[6:-1].split('-')
    else:
        residues = sequence.split('-')
    
    # Draw molecule first to get its bounds
    drawer.drawOptions().addAtomIndices = False
    drawer.DrawMolecule(mol)
    drawer.FinishDrawing()
    
    # Convert to PIL Image
    img = Image.open(BytesIO(drawer.GetDrawingText()))
    draw = ImageDraw.Draw(img)
    font = ImageFont.load_default(60)
    small_font = ImageFont.load_default(60)

    # Get molecule bounds
    conf = mol.GetConformer()
    positions = []
    for i in range(mol.GetNumAtoms()):
        pos = conf.GetAtomPosition(i)
        positions.append((pos.x, pos.y))
    
    x_coords = [p[0] for p in positions]
    y_coords = [p[1] for p in positions]
    min_x, max_x = min(x_coords), max(x_coords)
    min_y, max_y = min(y_coords), max(y_coords)
    
    # Calculate scaling factors
    scale = 150  # Increased scale factor
    center_x = 1000  # Image center
    center_y = 1000
    
    # Add residue labels in a circular arrangement around the structure
    n_residues = len(residues)
    radius = 700  # Distance of labels from center
    
    for i, residue in enumerate(residues):
        # Calculate position in a circle around the structure
        angle = (2 * np.pi * i / n_residues) - np.pi/2  # Start from top
        
        # Calculate label position
        label_x = center_x + radius * np.cos(angle)
        label_y = center_y + radius * np.sin(angle)
        
        # Draw residue label
        # Add white background for better visibility
        text = f"{i+1}. {residue}"
        bbox = draw.textbbox((label_x, label_y), text, font=font)
        padding = 10
        draw.rectangle([bbox[0]-padding, bbox[1]-padding, 
                       bbox[2]+padding, bbox[3]+padding], 
                      fill='white', outline='white')
        draw.text((label_x, label_y), text, 
                 font=font, fill='black', anchor="mm")
    
    # Add sequence at the top with white background
    seq_text = f"Sequence: {sequence}"
    bbox = draw.textbbox((center_x, 100), seq_text, font=small_font)
    padding = 10
    draw.rectangle([bbox[0]-padding, bbox[1]-padding, 
                   bbox[2]+padding, bbox[3]+padding], 
                  fill='white', outline='white')
    draw.text((center_x, 100), seq_text, 
             font=small_font, fill='black', anchor="mm")
    
    return img

def create_linear_peptide_viz(sequence):
    """
    Create a linear representation of peptide with residue annotations
    """
    # Create figure and axis
    fig, ax = plt.subplots(figsize=(15, 5))
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 2)
    
    # Parse sequence to get residues
    if sequence.startswith('cyclo('):
        residues = sequence[6:-1].split('-')  # Remove cyclo() and split
    else:
        residues = sequence.split('-')
    
    num_residues = len(residues)
    spacing = 9.0 / (num_residues - 1)  # Leave margins on sides
    
    # Draw peptide backbone
    y_pos = 1.5
    for i in range(num_residues):
        x_pos = 0.5 + i * spacing
        
        # Draw amino acid box
        rect = patches.Rectangle((x_pos-0.3, y_pos-0.2), 0.6, 0.4, 
                               facecolor='lightblue', edgecolor='black')
        ax.add_patch(rect)
        
        # Draw peptide bond
        if i < num_residues - 1:
            ax.plot([x_pos+0.3, x_pos+spacing-0.3], [y_pos, y_pos], 
                   color='black', linestyle='-', linewidth=2)
        
        # Add residue label with larger font
        ax.text(x_pos, y_pos-0.5, residues[i], ha='center', va='top', fontsize=14)
    
    # If cyclic, add arrow connecting ends
    if sequence.startswith('cyclo('):
        ax.annotate('', xy=(9.5, y_pos), xytext=(0.5, y_pos),
                   arrowprops=dict(arrowstyle='<->', color='red', lw=2))
        ax.text(5, y_pos+0.3, 'Cyclic Connection', ha='center', color='red', fontsize=14)
    
    # Add sequence at the top
    ax.text(5, 1.9, f"Sequence: {sequence}", ha='center', va='bottom', fontsize=12)
    
    # Remove axes
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis('off')
    
    return fig

def process_input(smiles_input=None, file_obj=None, show_linear=False):
    """Process input and create visualizations"""
    results = []
    images = []
    
    # Handle direct SMILES input
    if smiles_input:
        smiles = smiles_input.strip()
        
        # First check if it's a peptide
        if not is_peptide(smiles):
            return "Error: Input SMILES does not appear to be a peptide structure.", None, None
            
        try:
            # Create molecule
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Error: Invalid SMILES notation.", None, None
            
            # Get sequence and cyclic information
            sequence = parse_peptide(smiles)
            is_cyclic, peptide_cycles, aromatic_cycles = is_cyclic_peptide(smiles)
            
            # Create cyclic structure visualization
            img_cyclic = annotate_cyclic_structure(mol, sequence)
            
            # Create linear representation if requested
            img_linear = None
            if show_linear:
                fig_linear = create_linear_peptide_viz(sequence)
                
                # Convert matplotlib figure to image
                buf = BytesIO()
                fig_linear.savefig(buf, format='png', bbox_inches='tight', dpi=300)
                buf.seek(0)
                img_linear = Image.open(buf)
                plt.close(fig_linear)
            
            # Format text output
            output_text = f"Sequence: {sequence}\n"
            output_text += f"Is Cyclic: {'Yes' if is_cyclic else 'No'}\n"
            
            return output_text, img_cyclic, img_linear
            
        except Exception as e:
            return f"Error processing SMILES: {str(e)}", None, None
    
    # Handle file input
    if file_obj is not None:
        try:
            content = file_obj.decode('utf-8')
            output_text = ""
            for line in StringIO(content):
                smiles = line.strip()
                if smiles:
                    if not is_peptide(smiles):
                        output_text += f"Skipping non-peptide SMILES: {smiles}\n"
                        continue
                    result = analyze_single_smiles(smiles)
                    output_text += f"Sequence: {result['Sequence']}\n"
                    output_text += f"Is Cyclic: {result['Is Cyclic']}\n"
                    output_text += "-" * 50 + "\n"
            return output_text, None, None
        except Exception as e:
            return f"Error processing file: {str(e)}", None, None
    
    return "No input provided.", None, None

# Create Gradio interface
# [Previous imports and functions remain the same]

# Create Gradio interface with fixed examples
iface = gr.Interface(
    fn=process_input,
    inputs=[
        gr.Textbox(
            label="Enter SMILES string",
            placeholder="Enter SMILES notation of peptide...",
            lines=2
        ),
        gr.File(
            label="Or upload a text file with SMILES",
            file_types=[".txt"]
        ),
        gr.Checkbox(
            label="Show linear representation"
        )
    ],
    outputs=[
        gr.Textbox(
            label="Analysis Results",
            lines=10
        ),
        gr.Image(
            label="2D Structure with Annotations"
        ),
        gr.Image(
            label="Linear Representation"
        )
    ],
    title="Peptide Structure Analyzer and Visualizer",
    description="""
    Analyze and visualize peptide structures from SMILES notation:
    1. Validates if the input is a peptide structure
    2. Determines if the peptide is cyclic
    3. Parses the amino acid sequence
    4. Creates 2D structure visualization with residue annotations
    5. Optional linear representation
    
    Input: Either enter a SMILES string directly or upload a text file
    """,
    examples=[
        [
            "CC(C)C[C@@H]1NC(=O)[C@@H]2CCCN2C(=O)[C@@H](CC(C)C)NC(=O)[C@@H](CC(C)C)N(C)C(=O)[C@H](C)NC(=O)[C@H](Cc2ccccc2)NC1=O",
            None,  # Simply use None for file input in examples
            True
        ],
        [
            "CC(C)C[C@@H]1OC(=O)[C@H](C)NC(=O)[C@H](C(C)C)OC(=O)[C@H](C)N(C)C(=O)[C@@H](C)NC(=O)[C@@H](Cc2ccccc2)N(C)C1=O",
            None,
            True
        ]
    ],
    flagging_mode="never"
)

# Launch the app
if __name__ == "__main__":
    iface.launch()