add more functions

app.py

@@ -327,25 +327,44 @@ def annotate_cyclic_structure(mol, sequence):
     
     return img
 
-def create_linear_peptide_viz(sequence):
+def create_enhanced_linear_viz(sequence, smiles):
     """
-    Create a linear representation of peptide with residue annotations
+    Create an enhanced linear representation showing segment identification process
+    with improved segment handling
     """
-    # Create figure and axis
-    fig, ax = plt.subplots(figsize=(15, 5))
-    ax.set_xlim(0, 10)
-    ax.set_ylim(0, 2)
+    # 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 to get residues
+    # Parse sequence and get residues
     if sequence.startswith('cyclo('):
-        residues = sequence[6:-1].split('-')  # Remove cyclo() and split
+        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)  # Leave margins on sides
+    spacing = 9.0 / (num_residues - 1) if num_residues > 1 else 9.0
     
-    # Draw peptide backbone
+    # Draw basic structure
     y_pos = 1.5
     for i in range(num_residues):
         x_pos = 0.5 + i * spacing
@@ -353,30 +372,92 @@ def create_linear_peptide_viz(sequence):
         # 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)
+        ax_struct.add_patch(rect)
         
-        # Draw peptide bond
+        # Draw connecting bonds if not the last residue
         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)
+            # 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 residue label with larger font
-        ax.text(x_pos, y_pos-0.5, residues[i], ha='center', va='top', fontsize=14)
+        # 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 arrow connecting ends
+    # If cyclic, add connection indicator
     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)
+        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 sequence at the top
-    ax.text(5, 1.9, f"Sequence: {sequence}", ha='center', va='bottom', fontsize=12)
+    # 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
-    ax.set_xticks([])
-    ax.set_yticks([])
-    ax.axis('off')
+    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):
@@ -408,7 +489,7 @@ def process_input(smiles_input=None, file_obj=None, show_linear=False):
             # Create linear representation if requested
             img_linear = None
             if show_linear:
-                fig_linear = create_linear_peptide_viz(sequence)
+                fig_linear = create_enhanced_linear_viz(sequence, smiles)
                 
                 # Convert matplotlib figure to image
                 buf = BytesIO()