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 both natural and unnatural AAs
    Returns: tuple (residue, modifications)
    """
    modifications = []
    # Check for N-methylation
    if 'N(C)' in segment:  # Changed to look in current segment
        modifications.append('N-Me')
    if next_segment and 'OC(=O)' in next_segment:
        modifications.append('O-linked')

    # Check for Proline - but not if it's actually Cha
    if any(pattern in segment for pattern in ['CCCN2', 'N2CCC', '[C@@H]2CCCN2', 'CCCN1', 'N1CCC']):
        if not 'CCCCC' in segment:  # Make sure it's not Cha
            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) and not 'CCCCC' in combined:
                return ('Pro', modifications)

    # Check for O-tBu modification FIRST
    if 'COC(C)(C)C' in segment:
        return ('O-tBu', modifications)  # or return ('Ser(O-tBu)', modifications) if you prefer
        
    # Cyclohexyl amino acid (Cha)
    if 'N2CCCCC2' in segment or 'CCCCC2' in segment:
        return ('Cha', 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 - make Ala check more specific
    if '[C@H](CO)' in segment:
        return ('Ser', modifications)
    if '[C@@H]([C@@H](C)O)' in segment or '[C@H]([C@H](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)
    if ('[C@@H](C)' in segment or '[C@H](C)' in segment) and \
        not any(pat in segment for pat in ['O', 'CC(C)', 'COC']):
         return ('Ala', modifications)

 
    return (None, modifications)

def parse_peptide(smiles):
    """
    Parse peptide sequence with better segment identification
    """
    # Split at each peptide bond C(=O)N
    segments = []
    bonds = list(re.finditer(r'C\(=O\)N(?:\(C\))?', smiles))
    
    # Handle first residue (before first bond)
    first_bond = bonds[0].start()
    first_segment = smiles[0:first_bond]
    segments.append(first_segment)
    
    # Handle middle residues
    for i in range(len(bonds)):
        start = bonds[i].end()
        end = bonds[i+1].start() if i < len(bonds)-1 else len(smiles)
        segment = smiles[start:end]
        is_n_me = 'N(C)' in bonds[i].group()
        segments.append((segment, is_n_me))
    
    sequence = []
    # Handle first residue
    residue, mods = identify_residue(segments[0])
    if residue:
        sequence.append(residue)
    
    # Handle rest of residues
    for segment, is_n_me in segments[1:]:
        residue, mods = identify_residue(segment)
        if is_n_me:
            mods.append('N-Me')
        if residue:
            if mods:
                sequence.append(f"{residue}({','.join(mods)})")
            else:
                sequence.append(residue)
    
    print("\nDetailed Analysis:")
    print("Segments:", segments)
    print("Found sequence:", sequence)
    
    if is_cyclic_peptide(smiles):
        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'(\d)', smiles):
        number = match.group(1)
        position = match.start(1)
        
        if number not in cycle_info:
            cycle_info[number] = []
        cycle_info[number].append({
            'position': position,
            'full_context': smiles[max(0, position-3):min(len(smiles), position+4)]
        })
    
    # Print cycle information for debugging
    print("\nCycle Analysis:")
    for num, occurrences in cycle_info.items():
        print(f"Cycle number {num}:")
        for occ in occurrences:
            print(f"Position: {occ['position']}")
            print(f"Context: {occ['full_context']}")
    
    # Check each cycle
    peptide_cycles = []
    aromatic_cycles = []
    
    for number, occurrences in cycle_info.items():
        if len(occurrences) != 2:
            continue
            
        start, end = occurrences[0]['position'], occurrences[1]['position']
        
        # Get wider context for cycle classification
        segment = smiles[start:end+1]
        
        # First check if this is clearly an aromatic ring (phenylalanine side chain)
        full_context = smiles[max(0,start-10):min(len(smiles),end+10)]
        is_aromatic = ('c2ccccc2' in full_context and len(segment) < 20) or ('c1ccccc1' in full_context and len(segment) < 20)
        
        # Check for peptide bonds, including N-methylated ones
        peptide_patterns = [
            'C(=O)N',  # Regular peptide bond
            'C(=O)N(C)',  # N-methylated peptide bond
            'C(=O)N1',  # Cyclic peptide bond
            'C(=O)N2'   # Cyclic peptide bond
        ]
        
        # A peptide cycle should have multiple C(=O)N patterns and be longer
        has_peptide_bond = any(pattern in segment for pattern in peptide_patterns) and len(segment) > 20
        
        if is_aromatic and len(segment) < 20:  # Aromatic rings are typically shorter segments
            aromatic_cycles.append(number)
        elif has_peptide_bond:
            peptide_cycles.append(number)
    
    print("\nFound cycles:")
    print(f"Peptide cycles: {peptide_cycles}")
    print(f"Aromatic cycles: {aromatic_cycles}")
    
    return len(peptide_cycles) > 0

def analyze_single_smiles(smiles):
    """Analyze a single SMILES string"""
    try:
        is_cyclic, peptide_cycles, aromatic_cycles = is_cyclic_peptide(smiles)
        sidenote =  None
        sequence = parse_peptide(smiles)
        if is_cyclic and len(sequence) == 7:
            sidenote = 'This is some peptide sequence with modified side chains.'

        details = {
            #'SMILES': smiles,
            'Sequence': sequence,
            'Is Cyclic': 'Yes' if is_cyclic else 'No',
            'Sidenote': sidenote
            #'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
    # 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 and reverse it to match structural representation
    if sequence.startswith('cyclo('):
        residues = sequence[6:-1].split('-')
    else:
        residues = sequence.split('-')
    residues = list(reversed(residues))  # Reverse the sequence
    
    # 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)
    
    try:
        # Try to use DejaVuSans as it's commonly available on Linux systems
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 60)
        small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 60)
    except OSError:
        try:
            # Fallback to Arial if available (common on Windows)
            font = ImageFont.truetype("arial.ttf", 60)
            small_font = ImageFont.truetype("arial.ttf", 60)
        except OSError:
            # If no TrueType fonts are available, fall back to default
            print("Warning: TrueType fonts not available, using default font")
            font = ImageFont.load_default()
            small_font = ImageFont.load_default()
    # 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
    
    # Start from the rightmost point (3 o'clock position) and go counterclockwise
    # Offset by -3 positions to align with structure
    offset = 0  # Adjust this value to match the structure alignment
    for i, residue in enumerate(residues):
        # Calculate position in a circle around the structure
        # Start from 0 (3 o'clock) and go counterclockwise
        angle = -(2 * np.pi * ((i + offset) % n_residues) / n_residues)
        
        # Calculate label position
        label_x = center_x + radius * np.cos(angle)
        label_y = center_y + radius * np.sin(angle)
        
        # Draw residue label
        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 annotate_cyclic_structure(mol, sequence):
    """Create structure visualization with just the sequence header"""
    # Generate 2D coordinates
    AllChem.Compute2DCoords(mol)
    
    # Create drawer with larger size for annotations
    drawer = Draw.rdMolDraw2D.MolDraw2DCairo(2000, 2000)  # Even larger size
    
    # Draw molecule first
    drawer.drawOptions().addAtomIndices = False
    drawer.DrawMolecule(mol)
    drawer.FinishDrawing()
    
    # Convert to PIL Image
    img = Image.open(BytesIO(drawer.GetDrawingText()))
    draw = ImageDraw.Draw(img)
    try:
        small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 60)
    except OSError:
        try:
            small_font = ImageFont.truetype("arial.ttf", 60)
        except OSError:
            # If no TrueType fonts are available, fall back to default
            print("Warning: TrueType fonts not available, using default font")
            small_font = ImageFont.load_default()
    
    # Add just the sequence header at the top
    seq_text = f"Sequence: {sequence}"
    bbox = draw.textbbox((1000, 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((1000, 100), seq_text, 
             font=small_font, fill='black', anchor="mm")
    
    return img
def create_enhanced_linear_viz(sequence, smiles):
    """
    Create an enhanced linear representation showing segment identification process
    with improved segment handling
    """
    # Create figure with two subplots
    fig = plt.figure(figsize=(15, 10))
    gs = fig.add_gridspec(2, 1, height_ratios=[1, 2])
    ax_struct = fig.add_subplot(gs[0])
    ax_detail = fig.add_subplot(gs[1])
    
    # Parse sequence and get residues
    if sequence.startswith('cyclo('):
        residues = sequence[6:-1].split('-')
    else:
        residues = sequence.split('-')
    
    # Get molecule and analyze bonds
    mol = Chem.MolFromSmiles(smiles)
    
    # Split SMILES into segments for analysis
    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]  # Remove empty segments
    
    # Debug print
    print(f"Number of residues: {len(residues)}")
    print(f"Number of segments: {len(segments)}")
    print("Segments:", segments)
    
    # Top subplot - Basic structure
    ax_struct.set_xlim(0, 10)
    ax_struct.set_ylim(0, 2)
    
    num_residues = len(residues)
    spacing = 9.0 / (num_residues - 1) if num_residues > 1 else 9.0
    
    # Draw basic structure
    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_struct.add_patch(rect)
        
        # Draw connecting bonds if not the last residue
        if i < num_residues - 1:
            # Find the next bond pattern after this residue
            bond_segment = None
            for j in range(len(segments)):
                if re.match(bond_pattern, segments[j]):
                    if j > i*2 and j//2 == i:  # Found the right bond
                        bond_segment = segments[j]
                        break
            
            if bond_segment:
                bond_type, is_n_methylated = identify_linkage_type(bond_segment)
            else:
                bond_type = 'peptide'  # Default if not found
                
            bond_color = 'black' if bond_type == 'peptide' else 'red'
            linestyle = '-' if bond_type == 'peptide' else '--'
            
            # Draw bond line
            ax_struct.plot([x_pos+0.3, x_pos+spacing-0.3], [y_pos, y_pos], 
                         color=bond_color, linestyle=linestyle, linewidth=2)
            
            # Add bond type label
            mid_x = x_pos + spacing/2
            bond_label = f"{bond_type}"
            if is_n_methylated:
                bond_label += "\n(N-Me)"
            ax_struct.text(mid_x, y_pos+0.1, bond_label, 
                         ha='center', va='bottom', fontsize=10, 
                         color=bond_color)
        
        # Add residue label
        ax_struct.text(x_pos, y_pos-0.5, residues[i], 
                      ha='center', va='top', fontsize=14)
    
    # Bottom subplot - Detailed breakdown
    ax_detail.set_ylim(0, len(segments)+1)
    ax_detail.set_xlim(0, 1)
    
    # Create detailed breakdown
    segment_y = len(segments)  # Start from top
    for i, segment in enumerate(segments):
        y = segment_y - i
        
        # Check if this is a bond segment
        if re.match(bond_pattern, segment):
            bond_type, is_n_methylated = identify_linkage_type(segment)
            text = f"Bond {i//2 + 1}: {bond_type}"
            if is_n_methylated:
                text += " (N-methylated)"
            color = 'red'
        else:
            # Get next and previous segments for context
            next_seg = segments[i+1] if i+1 < len(segments) else None
            prev_seg = segments[i-1] if i > 0 else None
            
            residue, modifications = identify_residue(segment, next_seg, prev_seg)
            text = f"Residue {i//2 + 1}: {residue}"
            if modifications:
                text += f" ({', '.join(modifications)})"
            color = 'blue'
        
        # Add segment analysis
        ax_detail.text(0.05, y, text, fontsize=12, color=color)
        ax_detail.text(0.5, y, f"SMILES: {segment}", fontsize=10, color='gray')
    
    # If cyclic, add connection indicator
    if sequence.startswith('cyclo('):
        ax_struct.annotate('', xy=(9.5, y_pos), xytext=(0.5, y_pos),
                          arrowprops=dict(arrowstyle='<->', color='red', lw=2))
        ax_struct.text(5, y_pos+0.3, 'Cyclic Connection', 
                      ha='center', color='red', fontsize=14)
    
    # Add titles and adjust layout
    ax_struct.set_title("Peptide Structure Overview", pad=20)
    ax_detail.set_title("Segment Analysis Breakdown", pad=20)
    
    # Remove axes
    for ax in [ax_struct, ax_detail]:
        ax.set_xticks([])
        ax.set_yticks([])
        ax.axis('off')
    
    plt.tight_layout()
    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_enhanced_linear_viz(sequence, smiles)
                
                # 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:
            # Handle file content based on file object type
            if hasattr(file_obj, 'name'):  # If it's a file path
                with open(file_obj.name, 'r') as f:
                    content = f.read()
            else:  # If it's file content
                content = file_obj.decode('utf-8') if isinstance(file_obj, bytes) else str(file_obj)
            
            output_text = ""
            for line in content.splitlines():
                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 with simplified 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"],
            type="binary"
        ),
        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 containing SMILES strings
    
    Example SMILES strings (copy and paste):
    ```
    CC(C)C[C@@H]1NC(=O)[C@@H](CC(C)C)N(C)C(=O)[C@@H](C)N(C)C(=O)[C@H](Cc2ccccc2)NC(=O)[C@H](CC(C)C)N(C)C(=O)[C@H]2CCCN2C1=O
    ```
    ```
    C(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
    ```
    ```
    CC(C)C[C@H]1C(=O)N(C)[C@@H](Cc2ccccc2)C(=O)NCC(=O)N[C@H](C(=O)N2CCCCC2)CC(=O)N(C)CC(=O)N[C@@H]([C@@H](C)O)C(=O)N(C)[C@@H](C)C(=O)N[C@@H](COC(C)(C)C)C(=O)N(C)[C@@H](Cc2ccccc2)C(=O)N1C
    ```
    """,
    flagging_mode="never"
)

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