import os
import gradio as gr
import gradio.blocks
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
import tempfile
from rdkit import Chem

class PeptideAnalyzer:
    def __init__(self):
        self.bond_patterns = [
            #(r'OC\(=O\)', 'ester'),  # Ester bond
            (r'N\(C\)C\(=O\)', 'n_methyl'),  # N-methylated peptide bond
            (r'N[0-9]C\(=O\)', 'proline'),  # Proline peptide bond
            (r'NC\(=O\)', 'peptide'),  # Standard peptide bond
            (r'C\(=O\)N\(C\)', 'n_methyl_reverse'),  # Reverse N-methylated
            (r'C\(=O\)N[12]?', 'peptide_reverse')  # Reverse peptide bond
        ]
        self.complex_residue_patterns = [
            # Kpg - Lys(palmitoyl-Glu-OtBu) - Exact pattern for the specific structure
            (r'\[C[@]H\]\(CCCNC\(=O\)CCC\[C@@H\]\(NC\(=O\)CCCCCCCCCCCCCCCC\)C\(=O\)OC\(C\)\(C\)C\)', 'Kpg'),
            (r'CCCCCCCCCCCCCCCCC\(=O\)N\[C@H\]\(CCCC\(=O\)NCCC\[C@@H\]', 'Kpg'),
            (r'\[C@*H\]\(CSC\(c\d+ccccc\d+\)\(c\d+ccccc\d+\)c\d+ccc\(OC\)cc\d+\)', 'Cmt'),
            (r'CSC\(c.*?c.*?OC\)', 'Cmt'),        # Core structure of Cys-Mmt group
            (r'COc.*?ccc\(C\(SC', 'Cmt'),         # Start of Cmt in cyclic peptides
            (r'c2ccccc2\)c2ccccc2\)cc', 'Cmt'),   # End of Cmt in cyclic peptides
            # Glu(OAll) - Only match the complete pattern to avoid partial matches
            (r'C=CCOC\(=O\)CC\[C@@H\]', 'Eal'),
            (r'\(C\)OP\(=O\)\(O\)OCc\d+ccccc\d+', 'Tpb'),
            #(r'COc\d+ccc\(C\(SC\[C@@H\]\d+.*?\)\(c\d+ccccc\d+\)c\d+ccccc\d+\)cc\d+', 'Cmt-cyclic'),

            # Dtg - Asp(OtBu)-(Dmb)Gly - Full pattern
            (r'CN\(Cc\d+ccc\(OC\)cc\d+OC\)C\(=O\)\[C@H\]\(CC\(=O\)OC\(C\)\(C\)C\)', 'Dtg'),
            (r'C\(=O\)N\(CC\d+=C\(C=C\(C=C\d+\)OC\)OC\)CC\(=O\)', 'Dtg'),
            (r'N\[C@@H\]\(CC\(=O\)OC\(C\)\(C\)C\)C\(=O\)N\(CC\d+=C\(C=C\(C=C\d+\)OC\)OC\)CC\(=O\)', 'Dtg'),
        ]
        # Three to one letter code mapping
        self.three_to_one = {
            'Ala': 'A', 'Cys': 'C', 'Asp': 'D', 'Glu': 'E',
            'Phe': 'F', 'Gly': 'G', 'His': 'H', 'Ile': 'I',
            'Lys': 'K', 'Leu': 'L', 'Met': 'M', 'Asn': 'N',
            'Pro': 'P', 'Gln': 'Q', 'Arg': 'R', 'Ser': 'S',
            'Thr': 'T', 'Val': 'V', 'Trp': 'W', 'Tyr': 'Y',
            'ala': 'a', 'cys': 'c', 'asp': 'd', 'glu': 'e',
            'phe': 'f', 'gly': 'g', 'his': 'h', 'ile': 'i',
            'lys': 'k', 'leu': 'l', 'met': 'm', 'asn': 'n',
            'pro': 'p', 'gln': 'q', 'arg': 'r', 'ser': 's',
            'thr': 't', 'val': 'v', 'trp': 'w', 'tyr': 'y', 'Cmt-cyclic': 'Ĉ',
            'Aib': 'Ŷ', 'Dtg': 'Ĝ', 'Cmt': 'Ĉ', 'Eal': 'Ė', 'Nml': "Ŀ", 'Nma': 'Ṃ',
            'Kpg': 'Ƙ', 'Tpb': 'Ṯ', 'Cyl': 'Ċ', 'Nle': 'Ł', 'Hph': 'Ĥ', 'Cys-Cys': 'CC', 'cys-cys': 'cc',
        }
    def preprocess_complex_residues(self, smiles):
        """Identify and protect complex residues with internal peptide bonds - improved to prevent overlaps"""
        # Create a mapping of positions to complex residue types
        complex_positions = []
        
        # Search for all complex residue patterns
        for pattern, residue_type in self.complex_residue_patterns:
            for match in re.finditer(pattern, smiles):
                # Only add if this position doesn't overlap with existing matches
                if not any(pos['start'] <= match.start() < pos['end'] or 
                          pos['start'] < match.end() <= pos['end'] for pos in complex_positions):
                    complex_positions.append({
                        'start': match.start(),
                        'end': match.end(),
                        'type': residue_type,
                        'pattern': match.group()
                    })
        
        # Sort by position (to handle potential overlapping matches)
        complex_positions.sort(key=lambda x: x['start'])
        
        # If no complex residues found, return original SMILES
        if not complex_positions:
            return smiles, []
        
        # Build a new SMILES string, protecting complex residues
        preprocessed_smiles = smiles
        offset = 0  # Track offset from replacements
        
        protected_residues = []
        
        for pos in complex_positions:
            # Adjust positions based on previous replacements
            start = pos['start'] + offset
            end = pos['end'] + offset
            
            # Extract the complex residue part
            complex_part = preprocessed_smiles[start:end]
            
            # Verify this is a complete residue (should have proper amino acid structure)
            if not ('[C@H]' in complex_part or '[C@@H]' in complex_part):
                continue  # Skip if not a proper amino acid structure
                
            # Create a placeholder for this complex residue
            placeholder = f"COMPLEX_RESIDUE_{len(protected_residues)}"
            
            # Replace the complex part with the placeholder
            preprocessed_smiles = preprocessed_smiles[:start] + placeholder + preprocessed_smiles[end:]
            
            # Track the offset change
            offset += len(placeholder) - (end - start)
            
            # Store the residue information
            protected_residues.append({
                'placeholder': placeholder,
                'type': pos['type'],
                'content': complex_part
            })
            
            #print(f"Protected {pos['type']}: {complex_part[:20]}... as {placeholder}")
        
        return preprocessed_smiles, protected_residues
    def split_on_bonds(self, smiles, protected_residues=None):
        """Split SMILES into segments based on peptide bonds, with improved handling of protected residues"""
        positions = []
        used = set()
        
        # First, handle protected complex residues if any
        if protected_residues:
            for residue in protected_residues:
                match = re.search(residue['placeholder'], smiles)
                if match:
                    positions.append({
                        'start': match.start(),
                        'end': match.end(),
                        'type': 'complex',
                        'pattern': residue['placeholder'],
                        'residue_type': residue['type'],
                        'content': residue['content']
                    })
                    used.update(range(match.start(), match.end()))
        
        # Find all peptide bonds
        bond_positions = []
        
        # Find Gly pattern first
        gly_pattern = r'NCC\(=O\)'
        for match in re.finditer(gly_pattern, smiles):
            if not any(p in range(match.start(), match.end()) for p in used):
                bond_positions.append({
                    'start': match.start(),
                    'end': match.end(),
                    'type': 'gly',
                    'pattern': match.group()
                })
                used.update(range(match.start(), match.end()))
        
        # Then find all other bonds
        for pattern, bond_type in self.bond_patterns:
            for match in re.finditer(pattern, smiles):
                if not any(p in range(match.start(), match.end()) for p in used):
                    bond_positions.append({
                        'start': match.start(),
                        'end': match.end(),
                        'type': bond_type,
                        'pattern': match.group()
                    })
                    used.update(range(match.start(), match.end()))
        
        # Sort all positions
        bond_positions.sort(key=lambda x: x['start'])
        
        # Combine complex residue positions and bond positions
        all_positions = positions + bond_positions
        all_positions.sort(key=lambda x: x['start'])
        
        # Create segments
        segments = []
        
        # First segment (if not starting with a bond or complex residue)
        if all_positions and all_positions[0]['start'] > 0:
            segments.append({
                'content': smiles[0:all_positions[0]['start']],
                'bond_after': all_positions[0]['pattern'] if all_positions[0]['type'] != 'complex' else None,
                'complex_after': all_positions[0]['pattern'] if all_positions[0]['type'] == 'complex' else None
            })
        
        # Process segments between positions
        for i in range(len(all_positions)-1):
            current = all_positions[i]
            next_pos = all_positions[i+1]
            
            # Handle complex residues
            if current['type'] == 'complex':
                segments.append({
                    'content': current['content'],
                    'bond_before': all_positions[i-1]['pattern'] if i > 0 and all_positions[i-1]['type'] != 'complex' else None,
                    'bond_after': next_pos['pattern'] if next_pos['type'] != 'complex' else None,
                    'complex_type': current['residue_type']
                })
            # Handle regular bonds
            elif current['type'] == 'gly':
                segments.append({
                    'content': 'NCC(=O)',
                    'bond_before': all_positions[i-1]['pattern'] if i > 0 and all_positions[i-1]['type'] != 'complex' else None,
                    'bond_after': next_pos['pattern'] if next_pos['type'] != 'complex' else None
                })
            else:
                # Only create segment if there's content between this bond and next position
                content = smiles[current['end']:next_pos['start']]
                if content and next_pos['type'] != 'complex':
                    segments.append({
                        'content': content,
                        'bond_before': current['pattern'],
                        'bond_after': next_pos['pattern'] if next_pos['type'] != 'complex' else None
                    })
        
        # Last segment
        if all_positions and all_positions[-1]['end'] < len(smiles):
            if all_positions[-1]['type'] == 'complex':
                segments.append({
                    'content': all_positions[-1]['content'],
                    'bond_before': all_positions[-2]['pattern'] if len(all_positions) > 1 and all_positions[-2]['type'] != 'complex' else None,
                    'complex_type': all_positions[-1]['residue_type']
                })
            else:
                segments.append({
                    'content': smiles[all_positions[-1]['end']:],
                    'bond_before': all_positions[-1]['pattern']
                })
        
        return segments
    def is_peptide(self, smiles):
        """Check if the SMILES represents a peptide structure"""
        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
        
        return False

    def is_cyclic(self, smiles):
        """Improved cyclic peptide detection"""
        # Check for C-terminal carboxyl
        if smiles.endswith('C(=O)O'):
            return False, [], []
            
        # Find all numbers used in ring closures
        ring_numbers = re.findall(r'(?:^|[^c])[0-9](?=[A-Z@\(\)])', smiles)
        
        # Find aromatic ring numbers
        aromatic_matches = re.findall(r'c[0-9](?:ccccc|c\[nH\]c)[0-9]', smiles)
        aromatic_cycles = []
        for match in aromatic_matches:
            numbers = re.findall(r'[0-9]', match)
            aromatic_cycles.extend(numbers)
        
        # Numbers that aren't part of aromatic rings are peptide cycles
        peptide_cycles = [n for n in ring_numbers if n not in aromatic_cycles]
        
        is_cyclic = len(peptide_cycles) > 0 and not smiles.endswith('C(=O)O')
        return is_cyclic, peptide_cycles, aromatic_cycles
    

    def clean_terminal_carboxyl(self, segment):
        """Remove C-terminal carboxyl only if it's the true terminus"""
        content = segment['content']
        
        # Only clean if:
        # 1. Contains C(=O)O
        # 2. No bond_after exists (meaning it's the last segment)
        if 'C(=O)O' in content and not segment.get('bond_after'):
            # Remove C(=O)O pattern regardless of position
            cleaned = re.sub(r'\(C\(=O\)O\)', '', content)
            # Remove any leftover empty parentheses
            cleaned = re.sub(r'\(\)', '', cleaned)
            return cleaned
        return content

    def identify_residue(self, segment):
        """Identify residue with Pro reconstruction"""
        # Only clean terminal carboxyl if this is the last segment
        if 'complex_type' in segment:
            return segment['complex_type'], []
        content = self.clean_terminal_carboxyl(segment)
        mods = self.get_modifications(segment)
        
        if content.startswith('COc1ccc(C(SC[C@@H]'):
            print("DIRECT MATCH: Found Cmt at beginning")
            return 'Cmt', mods

        # VERY EXPLICIT check for the last segment in your example
        if '[C@@H]3CCCN3C2=O)(c2ccccc2)c2ccccc2)cc' in content:
            print("DIRECT MATCH: Found Pro at end")
            return 'Pro', mods
        # === Original amino acid patterns ===
        # Eal - Glu(OAll) - Multiple patterns
        if 'CCC(=O)OCC=C' in content or 'CC(=O)OCC=C' in content or 'C=CCOC(=O)CC' in content:
            return 'Eal', mods
        # Proline (P) - flexible ring numbers
        if any([
            # Check for any ring number in bond patterns
            (segment.get('bond_after', '').startswith(f'N{n}C(=O)') and 'CCC' in content and 
            any(f'[C@@H]{n}' in content or f'[C@H]{n}' in content for n in '123456789'))
            for n in '123456789'
        ]) or any([(segment.get('bond_before', '').startswith(f'C(=O)N{n}') and 'CCC' in content and 
                any(f'CCC{n}' for n in '123456789'))
                for n in '123456789'
        ]) or any([
            # Check ending patterns with any ring number
            (f'CCCN{n}' in content and content.endswith('=O') and 
            any(f'[C@@H]{n}' in content or f'[C@H]{n}' in content for n in '123456789'))
            for n in '123456789'
        ]) or any([
            # Handle CCC[C@H]n patterns
            (content == f'CCC[C@H]{n}' and segment.get('bond_before', '').startswith(f'C(=O)N{n}')) or
            (content == f'CCC[C@@H]{n}' and segment.get('bond_before', '').startswith(f'C(=O)N{n}')) or
            # N-terminal Pro with any ring number
            (f'N{n}CCC[C@H]{n}' in content) or
            (f'N{n}CCC[C@@H]{n}' in content)
            for n in '123456789'
        ]):
            return 'Pro', mods

        # D-Proline (p)
        if ('N1[C@H](CCC1)' in content):
            return 'pro', mods

        # Tryptophan (W) - more specific indole pattern
        if re.search(r'c[0-9]c\[nH\]c[0-9]ccccc[0-9][0-9]', content) and \
        'c[nH]c' in content.replace(' ', ''):
            # Check stereochemistry for D/L
            if '[C@H](CC' in content:  # D-form
                return 'trp', mods
            return 'Trp', mods
        
        # Lysine (K) - both patterns
        if '[C@@H](CCCCN)' in content or '[C@H](CCCCN)' in content:
            # Check stereochemistry for D/L
            if '[C@H](CCCCN)' in content:  # D-form
                return 'lys', mods
            return 'Lys', mods
                
        # Arginine (R) - both patterns
        if '[C@@H](CCCNC(=N)N)' in content or '[C@H](CCCNC(=N)N)' in content:
            # Check stereochemistry for D/L
            if '[C@H](CCCNC(=N)N)' in content:  # D-form
                return 'arg', mods
            return 'Arg', mods
        
        # Regular residue identification
        if content == 'C' and segment.get('bond_before') and segment.get('bond_after'):
            # If it's surrounded by peptide bonds, it's almost certainly Gly
            if ('C(=O)N' in segment['bond_before'] or 'NC(=O)' in segment['bond_before'] or 'N(C)C(=O)' in segment['bond_before']) and \
               ('NC(=O)' in segment['bond_after'] or 'C(=O)N' in segment['bond_after'] or 'N(C)C(=O)' in segment['bond_after']):
                return 'Gly', mods
            
        # Case 2: Cyclic terminal glycine - typically contains 'CNC' with ring closure
        if 'CNC' in content and any(f'C{i}=' in content for i in range(1, 10)):
            return 'Gly', mods  # This will catch patterns like 'CNC1=O'
        if not segment.get('bond_before') and segment.get('bond_after'):
            if content == 'C' or content == 'NC':
                if ('NC(=O)' in segment['bond_after'] or 'C(=O)N' in segment['bond_after'] or 'N(C)C(=O)' in segment['bond_after']):
                    return 'Gly', mods
                
        # Leucine patterns (L/l)
        if 'CC(C)C[C@H]' in content or 'CC(C)C[C@@H]' in content or '[C@@H](CC(C)C)' in content or '[C@H](CC(C)C)' in content or (('N[C@H](CCC(C)C)' in content or 'N[C@@H](CCC(C)C)' in content) and segment.get('bond_before') is None):
            # Check stereochemistry for D/L
            if '[C@H](CC(C)C)' in content or 'CC(C)C[C@H]' in content:  # D-form
                return 'leu', mods
            return 'Leu', mods

        # Threonine patterns (T/t)
        if '[C@@H]([C@@H](C)O)' in content or '[C@H]([C@H](C)O)' in content or '[C@@H]([C@H](C)O)' in content or '[C@H]([C@@H](C)O)' in content:
            # Check both stereochemistry patterns
            if '[C@H]([C@@H](C)O)' in content:  # D-form
                return 'thr', mods
            return 'Thr', mods
        
        if re.search(r'\[C@H\]\(CCc\d+ccccc\d+\)', content) or re.search(r'\[C@@H\]\(CCc\d+ccccc\d+\)', content):
            return 'Hph', mods    
        
        # Phenylalanine patterns (F/f)
        if re.search(r'\[C@H\]\(Cc\d+ccccc\d+\)', content) or re.search(r'\[C@@H\]\(Cc\d+ccccc\d+\)', content):
            # Check stereochemistry for D/L
            if re.search(r'\[C@H\]\(Cc\d+ccccc\d+\)', content):  # D-form
                return 'phe', mods
            return 'Phe', mods

        if ('CC(C)[C@@H]' in content or 'CC(C)[C@H]' in content or 
            '[C@H](C(C)C)' in content or '[C@@H](C(C)C)' in content or
            'C(C)C[C@H]' in content or 'C(C)C[C@@H]' in content):
            
            # Make sure it's not leucine
            if not any(p in content for p in ['CC(C)C[C@H]', 'CC(C)C[C@@H]', 'CCC(=O)']):
                # Check stereochemistry
                if '[C@H]' in content and not '[C@@H]' in content:  # D-form
                    return 'val', mods
                return 'Val', mods
                    
        # Isoleucine patterns (I/i)
        # First check for various isoleucine patterns while excluding valine
        if (any(['CC[C@@H](C)' in content, '[C@@H](C)CC' in content, '[C@@H](CC)C' in content, 
                'C(C)C[C@@H]' in content, '[C@@H]([C@H](C)CC)' in content, '[C@H]([C@@H](C)CC)' in content,
                '[C@@H]([C@@H](C)CC)' in content, '[C@H]([C@H](C)CC)' in content,
                'C[C@H](CC)[C@@H]' in content, 'C[C@@H](CC)[C@H]' in content, 
                'C[C@H](CC)[C@H]' in content, 'C[C@@H](CC)[C@@H]' in content,
                'CC[C@H](C)[C@@H]' in content, 'CC[C@@H](C)[C@H]' in content,
                'CC[C@H](C)[C@H]' in content, 'CC[C@@H](C)[C@@H]' in content]) 
            and 'CC(C)C' not in content):  # Exclude valine pattern
            
            # Check stereochemistry for D/L forms
            if any(['[C@H]([C@@H](CC)C)' in content, '[C@H](CC)C' in content,
                    '[C@H]([C@@H](C)CC)' in content, '[C@H]([C@H](C)CC)' in content,
                    'C[C@@H](CC)[C@H]' in content, 'C[C@H](CC)[C@H]' in content,
                    'CC[C@@H](C)[C@H]' in content, 'CC[C@H](C)[C@H]' in content]):
                # D-form
                return 'ile', mods
            # All other stereochemistries are treated as L-form
            return 'Ile', mods
        # Tpb - Thr(PO(OBzl)OH) - Multiple patterns
        if re.search(r'\(C\)OP\(=O\)\(O\)OCc[0-9]ccccc[0-9]', content) or 'OP(=O)(O)OCC' in content:
            return 'Tpb', mods
        
        # Alanine patterns (A/a)
        if ('[C@H](C)' in content or '[C@@H](C)' in content):
            if not any(p in content for p in ['C(C)C', 'COC', 'CN(', 'C(C)O', 'CC[C@H]', 'CC[C@@H]']):
                # Check stereochemistry for D/L
                if '[C@H](C)' in content:  # D-form
                    return 'ala', mods
                return 'Ala', mods
        
        # Tyrosine patterns (Y/y)
        if re.search(r'Cc[0-9]ccc\(O\)cc[0-9]', content):
            # Check stereochemistry for D/L
            if '[C@H](Cc1ccc(O)cc1)' in content:  # D-form
                return 'tyr', mods
            return 'Tyr', mods

        # Serine patterns (S/s)
        if '[C@H](CO)' in content or '[C@@H](CO)' in content:
            if not ('C(C)O' in content or 'COC' in content):
                # Check stereochemistry for D/L
                if '[C@H](CO)' in content:  # D-form
                    return 'ser', mods
                return 'Ser', mods
            
        if 'CSSC' in content:
            # Check for various cysteine-cysteine bridge patterns
            if re.search(r'\[C@@H\].*CSSC.*\[C@@H\]', content) or re.search(r'\[C@H\].*CSSC.*\[C@H\]', content):
                if '[C@H]' in content and not '[C@@H]' in content:  # D-form
                    return 'cys-cys', mods
                return 'Cys-Cys', mods
                
            # Pattern for cysteine with N-terminal amine group
            if '[C@@H](N)CSSC' in content or '[C@H](N)CSSC' in content:
                if '[C@H](N)CSSC' in content:  # D-form
                    return 'cys-cys', mods
                return 'Cys-Cys', mods
                
            # Pattern for cysteine with C-terminal carboxyl
            if 'CSSC[C@@H](C(=O)O)' in content or 'CSSC[C@H](C(=O)O)' in content:
                if 'CSSC[C@H](C(=O)O)' in content:  # D-form
                    return 'cys-cys', mods
                return 'Cys-Cys', mods
            
        # Cysteine patterns (C/c)
        if '[C@H](CS)' in content or '[C@@H](CS)' in content:
            # Check stereochemistry for D/L
            if '[C@H](CS)' in content:  # D-form
                return 'cys', mods
            return 'Cys', mods
        
        # Methionine patterns (M/m)
        if ('CCSC' in content) or ("CSCC" in content):
            # Check stereochemistry for D/L
            if '[C@H](CCSC)' in content:  # D-form
                return 'met', mods
            elif '[C@H]' in content:
                return 'met', mods
            return 'Met', mods
                    
        # Glutamine patterns (Q/q)
        if (content == '[C@@H](CC' or content == '[C@H](CC' and segment.get('bond_before')=='C(=O)N' and segment.get('bond_after')=='C(=O)N') or ('CCC(=O)N' in content) or ('CCC(N)=O' in content):
            # Check stereochemistry for D/L
            if '[C@H](CCC(=O)N)' in content:  # D-form
                return 'gln', mods
            return 'Gln', mods
            
        # Asparagine patterns (N/n)
        if (content == '[C@@H](C' or content == '[C@H](C' and segment.get('bond_before')=='C(=O)N' and segment.get('bond_after')=='C(=O)N') or ('CC(=O)N' in content) or ('CCN(=O)' in content) or ('CC(N)=O' in content):
            # Check stereochemistry for D/L
            if '[C@H](CC(=O)N)' in content:  # D-form
                return 'asn', mods
            return 'Asn', mods

        # Glutamic acid patterns (E/e)
        if ('CCC(=O)O' in content):
            # Check stereochemistry for D/L
            if '[C@H](CCC(=O)O)' in content:  # D-form
                return 'glu', mods
            return 'Glu', mods
            
        # Aspartic acid patterns (D/d)
        if ('CC(=O)O' in content):
            # Check stereochemistry for D/L
            if '[C@H](CC(=O)O)' in content:  # D-form
                return 'asp', mods
            return 'Asp', mods
        
        if re.search(r'Cc\d+c\[nH\]cn\d+', content) or re.search(r'Cc\d+cnc\[nH\]\d+', content):
            # Check stereochemistry for D/L
            if '[C@H]' in content:  # D-form
                return 'his', mods
            return 'His', mods
        if 'C2(CCCC2)' in content or 'C1(CCCC1)' in content or re.search(r'C\d+\(CCCC\d+\)', content):
            return 'Cyl', mods
        
        if ('N[C@@H](CCCC)' in content or '[C@@H](CCCC)' in content or 'CCCC[C@@H]' in content or 
            'N[C@H](CCCC)' in content or '[C@H](CCCC)' in content) and 'CC(C)' not in content:
            return 'Nle', mods
        # Aib - alpha-aminoisobutyric acid (2-aminoisobutyric acid)
        # More flexible pattern detection
        if 'C(C)(C)(N)' in content:
            return 'Aib', mods
        
        # Partial Aib pattern but NOT part of t-butyl ester
        if 'C(C)(C)' in content and 'OC(C)(C)C' not in content:
            if (segment.get('bond_before') and segment.get('bond_after') and
                any(bond in segment['bond_before'] for bond in ['C(=O)N', 'NC(=O)', 'N(C)C(=O)']) and
                any(bond in segment['bond_after'] for bond in ['NC(=O)', 'C(=O)N', 'N(C)C(=O)'])):
                return 'Aib', mods
        
        # Dtg - Asp(OtBu)-(Dmb)Gly - Simplified pattern for better detection
        if 'CC(=O)OC(C)(C)C' in content and 'CC1=C(C=C(C=C1)OC)OC' in content:
            return 'Dtg', mods
    
        
        # Kpg - Lys(palmitoyl-Glu-OtBu) - Simplified pattern
        if 'CCCNC(=O)' in content and 'CCCCCCCCCCCC' in content:
            return 'Kpg', mods
        
        
        
        return None, mods

    def get_modifications(self, segment):
        """Get modifications based on bond types and segment content - fixed to avoid duplicates"""
        mods = []
        
        # Check for N-methylation in any form, but only add it once
        # Check both bonds and segment content for N-methylation patterns
        if ((segment.get('bond_after') and 
            ('N(C)' in segment['bond_after'] or segment['bond_after'].startswith('C(=O)N(C)'))) or
            ('N(C)C(=O)' in segment['content'] or 'N(C)C1=O' in segment['content']) or
            (segment['content'].endswith('N(C)C(=O)') or segment['content'].endswith('N(C)C1=O'))):
            mods.append('N-Me')
        
        # Check for O-linked modifications
        #if segment.get('bond_after') and 'OC(=O)' in segment['bond_after']:
            #mods.append('O-linked')
        
        return mods
    
    def analyze_structure(self, smiles):
        """Main analysis function with preprocessing for complex residues"""
        #print("\nAnalyzing structure:", smiles)

        # Pre-process to identify complex residues first
        preprocessed_smiles, protected_residues = self.preprocess_complex_residues(smiles)
        """
        if protected_residues:
            print(f"Identified {len(protected_residues)} complex residues during pre-processing")
            for i, residue in enumerate(protected_residues):
                print(f"Complex residue {i+1}: {residue['type']}")
        """
        
        # Check if it's cyclic
        is_cyclic, peptide_cycles, aromatic_cycles = self.is_cyclic(smiles)
        
        # Split into segments, respecting protected residues
        segments = self.split_on_bonds(preprocessed_smiles, protected_residues)
        
        #print("\nSegment Analysis:")
        sequence = []
        for i, segment in enumerate(segments):
            """
            print(f"\nSegment {i}:")
            print(f"Content: {segment.get('content', 'None')}")
            print(f"Bond before: {segment.get('bond_before', 'None')}")
            print(f"Bond after: {segment.get('bond_after', 'None')}")
            """
            residue, mods = self.identify_residue(segment)
            if residue:
                if mods:
                    sequence.append(f"{residue}({','.join(mods)})")
                else:
                    sequence.append(residue)

                #print(f"Identified as: {residue}")
                #print(f"Modifications: {mods}")
            else:
                print(f"Warning: Could not identify residue in segment: {segment.get('content', 'None')}")
        
        # Format the sequence
        three_letter = '-'.join(sequence)
        
        # Use the mapping to create one-letter code
        one_letter = ''.join(self.three_to_one.get(aa.split('(')[0], 'X') for aa in sequence)
        
        if is_cyclic:
            three_letter = f"cyclo({three_letter})"
            one_letter = f"cyclo({one_letter})"
        """
        print(f"\nFinal sequence: {three_letter}")
        print(f"One-letter code: {one_letter}")
        print(f"Is cyclic: {is_cyclic}")
        print(f"Peptide cycles: {peptide_cycles}")
        print(f"Aromatic cycles: {aromatic_cycles}")
        """
        return {
            'three_letter': three_letter,
            'one_letter': one_letter,
            'is_cyclic': is_cyclic,
            'residues': sequence
        }
        
def annotate_cyclic_structure(mol, sequence):
    """Create structure visualization"""
    AllChem.Compute2DCoords(mol)
    
    drawer = Draw.rdMolDraw2D.MolDraw2DCairo(2000, 2000)
    
    # 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:
            print("Warning: TrueType fonts not available, using default font")
            small_font = ImageFont.load_default()
    
    # Header
    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):
    """"Linear visualization"""
    analyzer = PeptideAnalyzer()
    
    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])
    
    if sequence.startswith('cyclo('):
        residues = sequence[6:-1].split('-')
    else:
        residues = sequence.split('-')
    
    segments = analyzer.split_on_bonds(smiles)
    
    print(f"Number of residues: {len(residues)}")
    print(f"Number of segments: {len(segments)}")
    
    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
    
    y_pos = 1.5
    for i in range(num_residues):
        x_pos = 0.5 + i * spacing
        
        rect = patches.Rectangle((x_pos-0.3, y_pos-0.2), 0.6, 0.4, 
                               facecolor='lightblue', edgecolor='black')
        ax_struct.add_patch(rect)
        
        if i < num_residues - 1:
            segment = segments[i] if i < len(segments) else None
            if segment:
                bond_type = 'ester' if 'O-linked' in segment.get('bond_after', '') else 'peptide'
                is_n_methylated = 'N-Me' in segment.get('bond_after', '')
                
                bond_color = 'red' if bond_type == 'ester' else 'black'
                linestyle = '--' if bond_type == 'ester' else '-'
                
                ax_struct.plot([x_pos+0.3, x_pos+spacing-0.3], [y_pos, y_pos], 
                             color=bond_color, linestyle=linestyle, linewidth=2)
                
                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)
        
        ax_struct.text(x_pos, y_pos-0.5, residues[i], 
                      ha='center', va='top', fontsize=14)
    
    ax_detail.set_ylim(0, len(segments)+1)
    ax_detail.set_xlim(0, 1)
    
    segment_y = len(segments)
    for i, segment in enumerate(segments):
        y = segment_y - i
        
        # Check if this is a bond or residue
        residue, mods = analyzer.identify_residue(segment)
        if residue:
            text = f"Residue {i+1}: {residue}"
            if mods:
                text += f" ({', '.join(mods)})"
            color = 'blue'
        else:
            # Must be a bond
            text = f"Bond {i}: "
            if 'O-linked' in segment.get('bond_after', ''):
                text += "ester"
            elif 'N-Me' in segment.get('bond_after', ''):
                text += "peptide (N-methylated)"
            else:
                text += "peptide"
            color = 'red'
        
        ax_detail.text(0.05, y, text, fontsize=12, color=color)
        ax_detail.text(0.5, y, f"SMILES: {segment.get('content', '')}", 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)
    
    ax_struct.set_title("Peptide Structure Overview", pad=20)
    ax_detail.set_title("Segment Analysis Breakdown", pad=20)
    
    for ax in [ax_struct, ax_detail]:
        ax.set_xticks([])
        ax.set_yticks([])
        ax.axis('off')
    
    plt.tight_layout()
    return fig

class PeptideStructureGenerator:
    """Generate 3D structures of peptides using different embedding methods"""
    
    @staticmethod
    def prepare_molecule(smiles):
        """Prepare molecule with proper hydrogen handling"""
        mol = Chem.MolFromSmiles(smiles, sanitize=False)
        if mol is None:
            raise ValueError("Failed to create molecule from SMILES")
        
        for atom in mol.GetAtoms():
            atom.UpdatePropertyCache(strict=False)
        
        # Sanitize with reduced requirements
        Chem.SanitizeMol(mol, 
                        sanitizeOps=Chem.SANITIZE_FINDRADICALS|
                                  Chem.SANITIZE_KEKULIZE|
                                  Chem.SANITIZE_SETAROMATICITY|
                                  Chem.SANITIZE_SETCONJUGATION|
                                  Chem.SANITIZE_SETHYBRIDIZATION|
                                  Chem.SANITIZE_CLEANUPCHIRALITY)
        
        mol = Chem.AddHs(mol)
        return mol

    @staticmethod
    def get_etkdg_params(attempt=0):
        """Get ETKDG parameters"""
        params = AllChem.ETKDGv3()
        params.randomSeed = -1
        params.maxIterations = 200
        params.numThreads = 4  # Reduced for web interface
        params.useBasicKnowledge = True
        params.enforceChirality = True
        params.useExpTorsionAnglePrefs = True
        params.useSmallRingTorsions = True
        params.useMacrocycleTorsions = True
        params.ETversion = 2
        params.pruneRmsThresh = -1
        params.embedRmsThresh = 0.5
        
        if attempt > 10:
            params.bondLength = 1.5 + (attempt - 10) * 0.02
            params.useExpTorsionAnglePrefs = False
        
        return params

    def generate_structure_etkdg(self, smiles, max_attempts=20):
        """Generate 3D structure using ETKDG without UFF optimization"""
        success = False
        mol = None
        
        for attempt in range(max_attempts):
            try:
                mol = self.prepare_molecule(smiles)
                params = self.get_etkdg_params(attempt)
                
                if AllChem.EmbedMolecule(mol, params) == 0:
                    success = True
                    break
            except Exception as e:
                continue
        
        if not success:
            raise ValueError("Failed to generate structure with ETKDG")
        
        return mol

    def generate_structure_uff(self, smiles, max_attempts=20):
        """Generate 3D structure using ETKDG followed by UFF optimization"""
        best_mol = None
        lowest_energy = float('inf')
        
        for attempt in range(max_attempts):
            try:
                test_mol = self.prepare_molecule(smiles)
                params = self.get_etkdg_params(attempt)
                
                if AllChem.EmbedMolecule(test_mol, params) == 0:
                    res = AllChem.UFFOptimizeMolecule(test_mol, maxIters=2000, 
                                                     vdwThresh=10.0, confId=0,
                                                     ignoreInterfragInteractions=True)
                    
                    if res == 0:
                        ff = AllChem.UFFGetMoleculeForceField(test_mol)
                        if ff:
                            current_energy = ff.CalcEnergy()
                            if current_energy < lowest_energy:
                                lowest_energy = current_energy
                                best_mol = Chem.Mol(test_mol)
            except Exception:
                continue
        
        if best_mol is None:
            raise ValueError("Failed to generate optimized structure")
        
        return best_mol

    @staticmethod
    def mol_to_sdf_bytes(mol):
        """Convert RDKit molecule to SDF file bytes"""
        sio = StringIO()
        writer = Chem.SDWriter(sio)
        writer.write(mol)
        writer.close()
        
        return sio.getvalue().encode('utf-8')

def process_input(
    smiles_input=None, 
    file_obj=None, 
    show_linear=False, 
    show_segment_details=False, 
    generate_3d=False, 
    use_uff=False
):
    """Process input and create visualizations using PeptideAnalyzer"""
    analyzer = PeptideAnalyzer()
    temp_dir = tempfile.mkdtemp() if generate_3d else None
    structure_files = []
    
    # Handle direct SMILES input
    if smiles_input:
        smiles = smiles_input.strip()
        
        if not analyzer.is_peptide(smiles):
            return "Error: Input SMILES does not appear to be a peptide structure.", None, None, []
                
        try:
            # Preprocess to protect complex residues
            pre_smiles, protected_residues = analyzer.preprocess_complex_residues(smiles)
            # Report protected residues in summary if any
            protected_info = None
            if protected_residues:
                protected_info = [res['type'] for res in protected_residues]
                
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Error: Invalid SMILES notation.", None, None, []
            
            if generate_3d:
                generator = PeptideStructureGenerator()
                
                try:
                    # Generate ETKDG structure
                    mol_etkdg = generator.generate_structure_etkdg(smiles)
                    etkdg_path = os.path.join(temp_dir, "structure_etkdg.sdf")
                    writer = Chem.SDWriter(etkdg_path)
                    writer.write(mol_etkdg)
                    writer.close()
                    structure_files.append(etkdg_path)
                    
                    # Generate UFF structure if requested
                    if use_uff:
                        mol_uff = generator.generate_structure_uff(smiles)
                        uff_path = os.path.join(temp_dir, "structure_uff.sdf")
                        writer = Chem.SDWriter(uff_path)
                        writer.write(mol_uff)
                        writer.close()
                        structure_files.append(uff_path)
                
                except Exception as e:
                    return f"Error generating 3D structures: {str(e)}", None, None, []
            
            analysis = analyzer.analyze_structure(smiles)
            three_letter = analysis['three_letter']
            one_letter = analysis['one_letter']
            is_cyclic = analysis['is_cyclic']

            # Only include segment analysis in output if requested
            if show_segment_details:
                segments = analyzer.split_on_bonds(smiles)
            
                sequence_parts = []
                output_text = ""
                output_text += "Segment Analysis:\n"
                for i, segment in enumerate(segments):
                    output_text += f"\nSegment {i}:\n"
                    output_text += f"Content: {segment['content']}\n"
                    output_text += f"Bond before: {segment.get('bond_before', 'None')}\n"
                    output_text += f"Bond after: {segment.get('bond_after', 'None')}\n"
                    
                    residue, mods = analyzer.identify_residue(segment)
                    if residue:
                        if mods:
                            sequence_parts.append(f"{residue}({','.join(mods)})")
                        else:
                            sequence_parts.append(residue)
                        output_text += f"Identified as: {residue}\n"
                        output_text += f"Modifications: {mods}\n"
                    else:
                        output_text += f"Warning: Could not identify residue in segment: {segment['content']}\n"
                output_text += "\n"
                is_cyclic, peptide_cycles, aromatic_cycles = analyzer.is_cyclic(smiles)
                three_letter = '-'.join(sequence_parts)
                one_letter = ''.join(analyzer.three_to_one.get(aa.split('(')[0], 'X') for aa in sequence_parts)
            else:
                pass
            
            img_cyclic = annotate_cyclic_structure(mol, three_letter)
            
            # Create linear representation if requested
            img_linear = None
            if show_linear:
                fig_linear = create_enhanced_linear_viz(three_letter, smiles)
                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)

            summary = "Summary:\n"
            summary += f"Sequence: {three_letter}\n"
            summary += f"One-letter code: {one_letter}\n"
            summary += f"Is Cyclic: {'Yes' if is_cyclic else 'No'}\n"
            #if is_cyclic:
                #summary += f"Peptide Cycles: {', '.join(peptide_cycles)}\n"
                #summary += f"Aromatic Cycles: {', '.join(aromatic_cycles)}\n"
            
            if structure_files:
                summary += "\n3D Structures Generated:\n"
                for filepath in structure_files:
                    summary += f"- {os.path.basename(filepath)}\n"
            
            #return summary, img_cyclic, img_linear, structure_files if structure_files else None
            return summary, img_cyclic

        except Exception as e:
            #return f"Error processing SMILES: {str(e)}", None, None, []
            return f"Error processing SMILES: {str(e)}", None
    # Handle file input
    if file_obj is not None:
        try:
            if hasattr(file_obj, 'name'):
                with open(file_obj.name, 'r') as f:
                    content = f.read()
            else:
                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 not smiles:
                    continue
                
                if not analyzer.is_peptide(smiles):
                    output_text += f"Skipping non-peptide SMILES: {smiles}\n"
                    continue
                
                try:
                    # Process the structure
                    result = analyzer.analyze_structure(smiles)
                    
                    output_text += f"\nSummary for SMILES: {smiles}\n"
                    output_text += f"Sequence: {result['three_letter']}\n"
                    output_text += f"One-letter code: {result['one_letter']}\n"
                    output_text += f"Is Cyclic: {'Yes' if result['is_cyclic'] else 'No'}\n"
                    output_text += "-" * 50 + "\n"
                except Exception as e:
                    output_text += f"Error processing SMILES: {smiles} - {str(e)}\n"
                    output_text += "-" * 50 + "\n"
            
            return output_text, None, None, []
            
        except Exception as e:
            return f"Error processing file: {str(e)}", None, None, []
    
    return (
            output_text or "No analysis done.",
            img_cyclic if 'img_cyclic' in locals() else None,
            #img_linear if 'img_linear' in locals() else None,
            #structure_files if structure_files else []
        )
"""
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"]
        #)],
    outputs=[
        gr.Textbox(
            label="Analysis Results",
            lines=10
        ),
        gr.Image(
            label="2D Structure with Annotations",
            type="pil"
        ),
    ],
    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
    6. Optional 3D structure generation (ETKDG and UFF methods)
    
    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"
)

if __name__ == "__main__":
    iface.launch(share=True)
"""
from fastapi import FastAPI
import gradio as gr

# 1) Make a FastAPI with no OpenAPI/docs routes
app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)

# 2) Build your Interface as before
iface = gr.Interface(
    fn=process_input,
    inputs=[ gr.Textbox(label="Enter SMILES string", lines=2) ],
    outputs=[
      gr.Textbox(label="Analysis Results", lines=10),
      gr.Image(label="2D Structure with Annotations", type="pil"),
    ],
    title="Peptide Structure Analyzer and Visualizer",
    flagging_mode="never"
)

# 3) Mount it at “/”
app = gr.mount_gradio_app(app, iface, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)