Update app.py

app.py

@@ -43,7 +43,7 @@ def apply_custom_styling():
         <style>
         @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@400;700&display=swap');
         
-        html, body, [class*=\"st-\"] {
+        html, body, [class*="st-"] {
             font-family: 'Roboto', sans-serif;
         }
 
@@ -53,11 +53,11 @@ def apply_custom_styling():
         }
 
         /* Tab styles */
-        .stTabs [data-baseweb=\"tab-list\"] {
+        .stTabs [data-baseweb="tab-list"] {
             gap: 24px;
         }
 
-        .stTabs [data-baseweb=\"tab\"] {
+        .stTabs [data-baseweb="tab"] {
             height: 50px;
             white-space: pre-wrap;
             background: none;
@@ -67,12 +67,12 @@ def apply_custom_styling():
             color: #AAA;
         }
         
-        .stTabs [data-baseweb=\"tab\"]:hover {
+        .stTabs [data-baseweb="tab"]:hover {
             background: #222;
             color: #FFF;
         }
 
-        .stTabs [aria-selected=\"true\"] {
+        .stTabs [aria-selected="true"] {
             border-bottom: 2px solid #00A0FF; /* Highlight color for active tab */
             color: #FFF;
         }
@@ -154,7 +154,6 @@ def visualize_protein_3d(pdb_data: str, title="Protein 3D Structure"):
         return None, "Cannot generate 3D view: No PDB data provided."
     try:
         viewer = py3Dmol.view(width='100%', height=600)
-        # MODIFIED: Changed background color
         viewer.setBackgroundColor('#1C1C1C') 
         viewer.addModel(pdb_data, "pdb")
         viewer.setStyle({'cartoon': {'color': 'spectrum', 'thickness': 0.8}})
@@ -206,28 +205,46 @@ def calculate_molecular_properties(smiles_list: list):
 
 def assess_drug_likeness(df: pd.DataFrame):
     """
-    Assesses drug-likeness based on Lipinski's Rule of Five. 
+    Assesses drug-likeness based on Lipinski's Rule of Five.
+    This version returns a boolean for plotting and a formatted string for display.
     """
     if df.empty:
-        return pd.DataFrame(), "Cannot assess drug-likeness: No properties data." [cite: 26]
-    df_copy = df.copy()
-    df_copy['MW_OK'] = df_copy['MW'] <= 500
-    df_copy['LogP_OK'] = df_copy['LogP'] <= 5
-    df_copy['HBD_OK'] = df_copy['HBD'] <= 5
-    df_copy['HBA_OK'] = df_copy['HBA'] <= 10
-    df_copy['Lipinski_Violations'] = (~df_copy[['MW_OK', 'LogP_OK', 'HBD_OK', 'HBA_OK']]).sum(axis=1)
-    # Fixed: Use proper colored emojis instead of boolean values
-    df_copy['Drug_Like'] = df_copy['Lipinski_Violations'].apply(lambda x: '✅ Yes' if x <= 1 else '❌ No')
-    log = "✅ Assessed drug-likeness using Lipinski's Rule of Five.\n" [cite: 27]
-    return df_copy, log [cite: 27]
+        return pd.DataFrame(), pd.DataFrame(), "Cannot assess drug-likeness: No properties data."
     
+    # Create a copy for analysis to avoid modifying the original dataframe
+    analysis_df = df.copy()
+    analysis_df['MW_OK'] = analysis_df['MW'] <= 500
+    analysis_df['LogP_OK'] = analysis_df['LogP'] <= 5
+    analysis_df['HBD_OK'] = analysis_df['HBD'] <= 5
+    analysis_df['HBA_OK'] = analysis_df['HBA'] <= 10
+    analysis_df['Lipinski_Violations'] = (~analysis_df[['MW_OK', 'LogP_OK', 'HBD_OK', 'HBA_OK']]).sum(axis=1)
+    
+    # This boolean column is for the plotting function
+    analysis_df['Drug_Like'] = analysis_df['Lipinski_Violations'] <= 1
+    
+    # Create a separate DataFrame for display purposes with emojis
+    display_df = df.copy()
+    display_df['Lipinski_Violations'] = analysis_df['Lipinski_Violations']
+    display_df['Drug_Like'] = analysis_df['Drug_Like'].apply(lambda x: '✅ Yes' if x else '❌ No')
+    
+    log = "✅ Assessed drug-likeness using Lipinski's Rule of Five.\n"
+    
+    # Return both the analysis_df (for plotting) and display_df (for tables)
+    return analysis_df, display_df, log
+
+
 def plot_properties_dashboard(df: pd.DataFrame):
     """
-    Creates a 2x2 dashboard of molecular property visualizations. 
+    Creates a 2x2 dashboard of molecular property visualizations.
+    This version expects a boolean 'Drug_Like' column.
     """
-    if df.empty:
-        return None, "Cannot plot: No analysis data." [cite: 28]
+    if df.empty or 'Drug_Like' not in df.columns:
+        return None, "Cannot plot: No analysis data or 'Drug_Like' column missing."
     
+    # Ensure 'Drug_Like' is boolean for correct mapping
+    if df['Drug_Like'].dtype != bool:
+         return None, f"Cannot plot: 'Drug_Like' column must be boolean, but it is {df['Drug_Like'].dtype}."
+
     plt.style.use('dark_background')
     fig, axes = plt.subplots(2, 2, figsize=(12, 10))
     fig.suptitle("Molecular Properties Analysis", fontsize=16)
@@ -237,17 +254,18 @@ def plot_properties_dashboard(df: pd.DataFrame):
         for ax in ax_row:
             ax.set_facecolor('none')
 
-    # FIXED: Correctly map string values to colors
-    axes[0,0].scatter(df['MW'], df['LogP'], c=df['Drug_Like'].map({'✅ Yes': 'green', '❌ No': 'red'}), s=80, alpha=0.7)
+    # Color mapping now correctly uses boolean values
+    color_map = {True: 'green', False: 'red'}
+    
+    axes[0,0].scatter(df['MW'], df['LogP'], c=df['Drug_Like'].map(color_map), s=80, alpha=0.7)
     axes[0,0].set_title('Molecular Weight vs LogP')
-    axes[0,0].set_xlabel('Molecular Weight (Da)') [cite: 29]
-    axes[0,0].set_ylabel('LogP') [cite: 29]
-    axes[0,0].axvline(500, color='r', linestyle='--', alpha=0.6, label='MW < 500') [cite: 29]
-    axes[0,0].axhline(5, color='r', linestyle='--', alpha=0.6, label='LogP < 5') [cite: 29]
+    axes[0,0].set_xlabel('Molecular Weight (Da)')
+    axes[0,0].set_ylabel('LogP')
+    axes[0,0].axvline(500, color='r', linestyle='--', alpha=0.6, label='MW < 500')
+    axes[0,0].axhline(5, color='r', linestyle='--', alpha=0.6, label='LogP < 5')
     axes[0,0].legend()
 
-    # FIXED: Correctly map string values to colors
-    axes[0,1].scatter(df['HBD'], df['HBA'], c=df['Drug_Like'].map({'✅ Yes': 'green', '❌ No': 'red'}), s=80, alpha=0.7)
+    axes[0,1].scatter(df['HBD'], df['HBA'], c=df['Drug_Like'].map(color_map), s=80, alpha=0.7)
     axes[0,1].set_title('Hydrogen Bonding Properties')
     axes[0,1].set_xlabel('Hydrogen Bond Donors')
     axes[0,1].set_ylabel('Hydrogen Bond Acceptors')
@@ -255,22 +273,20 @@ def plot_properties_dashboard(df: pd.DataFrame):
     axes[0,1].axhline(10, color='r', linestyle='--', alpha=0.6, label='HBA < 10')
     axes[0,1].legend()
 
-    # FIXED: Correctly map string values to colors
-    axes[1,0].scatter(df['TPSA'], df['RotBonds'], c=df['Drug_Like'].map({'✅ Yes': 'green', '❌ No': 'red'}), s=80, alpha=0.7)
-    axes[1,0].set_title('TPSA vs Flexibility') [cite: 30]
-    axes[1,0].set_xlabel('Topological Polar Surface Area (Ų)') [cite: 30]
-    axes[1,0].set_ylabel('Rotatable Bonds') [cite: 30]
+    axes[1,0].scatter(df['TPSA'], df['RotBonds'], c=df['Drug_Like'].map(color_map), s=80, alpha=0.7)
+    axes[1,0].set_title('TPSA vs Flexibility')
+    axes[1,0].set_xlabel('Topological Polar Surface Area (Ų)')
+    axes[1,0].set_ylabel('Rotatable Bonds')
 
     drug_like_counts = df['Drug_Like'].value_counts()
-    # FIXED: Correctly map string values to labels and colors
-    labels = drug_like_counts.index
-    colors = ['green' if i == '✅ Yes' else 'red' for i in drug_like_counts.index]
+    labels = ['Drug-like' if i else 'Non-drug-like' for i in drug_like_counts.index]
+    colors = ['green' if i else 'red' for i in drug_like_counts.index]
     axes[1,1].pie(drug_like_counts.values, labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
     axes[1,1].set_title('Drug-likeness Distribution')
     
     plt.tight_layout(rect=[0, 0, 1, 0.96])
     return fig, "✅ Generated properties dashboard."
-    
+
 # ===== Phase 2 Functions =====
 def get_phase2_molecules():
     return {
@@ -312,21 +328,16 @@ def visualize_molecule_2d_3d(smiles: str, name: str):
         if not mol: return f"<p>Invalid SMILES for {name}</p>", f"❌ Invalid SMILES for {name}"
         
         drawer = Draw.rdMolDraw2D.MolDraw2DSVG(400, 300)
-        # Set dark theme colors for 2D drawing
         drawer.drawOptions().clearBackground = False
         drawer.drawOptions().addStereoAnnotation = True
         drawer.drawOptions().baseFontSize = 0.8
         drawer.drawOptions().circleAtoms = False
-        drawer.drawOptions().highlightColour = (1, 0.5, 0)  # Orange for highlights
-        
-        # Set colors for dark background visibility
-        drawer.drawOptions().backgroundColour = (0.11, 0.11, 0.11)  # Dark background
-        drawer.drawOptions().symbolColour = (1, 1, 1)  # White symbols
-        drawer.drawOptions().defaultColour = (1, 1, 1)  # White default color
-        
-        # Try to set annotation color (this might help with (R)/(S) labels)
+        drawer.drawOptions().highlightColour = (1, 0.5, 0)
+        drawer.drawOptions().backgroundColour = (0.11, 0.11, 0.11)
+        drawer.drawOptions().symbolColour = (1, 1, 1)
+        drawer.drawOptions().defaultColour = (1, 1, 1)
         try:
-            drawer.drawOptions().annotationColour = (1, 1, 1)  # White annotations
+            drawer.drawOptions().annotationColour = (1, 1, 1)
         except:
             pass
 
@@ -334,52 +345,13 @@ def visualize_molecule_2d_3d(smiles: str, name: str):
         drawer.FinishDrawing()
         svg_2d = drawer.GetDrawingText().replace('svg:', '')
 
-        # More aggressive SVG text color fixes - target all possible black text variations
         import re
-        
-        # First, comprehensive string replacements
         svg_2d = svg_2d.replace('stroke="black"', 'stroke="white"')
         svg_2d = svg_2d.replace('fill="black"', 'fill="white"')
         svg_2d = svg_2d.replace('stroke="#000000"', 'stroke="#FFFFFF"')
         svg_2d = svg_2d.replace('fill="#000000"', 'fill="#FFFFFF"')
-        svg_2d = svg_2d.replace('stroke="#000"', 'stroke="#FFF"')
-        svg_2d = svg_2d.replace('fill="#000"', 'fill="#FFF"')
-        svg_2d = svg_2d.replace('stroke:black', 'stroke:white')
-        svg_2d = svg_2d.replace('fill:black', 'fill:white')
-        svg_2d = svg_2d.replace('stroke:#000000', 'stroke:#FFFFFF')
-        svg_2d = svg_2d.replace('fill:#000000', 'fill:#FFFFFF')
-        svg_2d = svg_2d.replace('stroke:#000', 'stroke:#FFF')
-        svg_2d = svg_2d.replace('fill:#000', 'fill:#FFF')
-        svg_2d = svg_2d.replace('stroke="rgb(0,0,0)"', 'stroke="rgb(255,255,255)"')
-        svg_2d = svg_2d.replace('fill="rgb(0,0,0)"', 'fill="rgb(255,255,255)"')
-        svg_2d = svg_2d.replace('stroke:rgb(0,0,0)', 'stroke:rgb(255,255,255)')
-        svg_2d = svg_2d.replace('fill:rgb(0,0,0)', 'fill:rgb(255,255,255)')
-        svg_2d = svg_2d.replace('color="black"', 'color="white"')
-        svg_2d = svg_2d.replace('color:#000000', 'color:#FFFFFF')
-        svg_2d = svg_2d.replace('color:#000', 'color:#FFF')
-        
-        # Aggressive regex-based fixes for all text elements
-        # Remove any existing fill attributes from text elements and add white fill
-        svg_2d = re.sub(r'<text([^>]*?)\s+fill="[^"]*"([^>]*?)>', r'<text\1\2 fill="white">', svg_2d)
-        svg_2d = re.sub(r'<text([^>]*?)(?<!fill="white")>', r'<text\1 fill="white">', svg_2d)
-        
-        # Fix style attributes in text elements
-        svg_2d = re.sub(r'<text([^>]*?)style="([^"]*?)fill:\s*(?:black|#000000|#000|rgb\(0,0,0\))([^"]*?)"([^>]*?)>', 
-                       r'<text\1style="\2fill:white\3"\4>', svg_2d)
-        
-        # If text elements don't have any fill specified, ensure they get white
         svg_2d = re.sub(r'<text(?![^>]*fill=)([^>]*?)>', r'<text fill="white"\1>', svg_2d)
         
-        # Clean up any duplicate fill attributes
-        svg_2d = re.sub(r'fill="white"\s+fill="white"', 'fill="white"', svg_2d)
-        
-        # Final catch-all: replace any remaining black in the entire SVG
-        svg_2d = re.sub(r'\bblack\b', 'white', svg_2d)
-        svg_2d = re.sub(r'#000000', '#FFFFFF', svg_2d)
-        svg_2d = re.sub(r'#000\b', '#FFF', svg_2d)
-        svg_2d = re.sub(r'rgb\(0,\s*0,\s*0\)', 'rgb(255,255,255)', svg_2d)
-        
-        # Embed the SVG within a div with a dark background for consistency
         svg_2d = f'<div style="background-color: #1C1C1C; padding: 10px; border-radius: 5px;">{svg_2d}</div>'
 
         mol_3d = Chem.AddHs(mol)
@@ -416,7 +388,6 @@ def visualize_protein_ligand_interaction(pdb_data: str, pdb_id: str, ligand_resn
     if not pdb_data: return None, "Cannot generate view: No PDB data provided."
     try:
         viewer = py3Dmol.view(width='100%', height=700)
-        # MODIFIED: Changed background color
         viewer.setBackgroundColor('#1C1C1C')
         viewer.addModel(pdb_data, "pdb")
         viewer.setStyle({'cartoon': {'color': 'spectrum', 'thickness': 0.8}})
@@ -453,7 +424,6 @@ def calculate_comprehensive_properties(smiles_dict: dict):
                          'TPSA': Descriptors.TPSA(mol), 'Rotatable_Bonds': Descriptors.NumRotatableBonds(mol),
                          'Aromatic_Rings': Descriptors.NumAromaticRings(mol),
                          'Lipinski_Violations': violations, 
-                         # Fixed: Use proper colored emojis instead of text
                          'Drug_Like': '✅ Yes' if violations <= 1 else '❌ No'})
     df = pd.DataFrame(analysis).round(2)
     log += f"✅ Calculated comprehensive properties for {len(df)} compounds.\n"
@@ -478,7 +448,6 @@ def predict_toxicity(properties_df: pd.DataFrame):
     toxicity_prob = rf_model.predict_proba(X_pred)[:, 1]
     results_df = properties_df[['Compound']].copy()
     results_df['Toxicity_Probability'] = np.round(toxicity_prob, 3)
-    # Fixed: Use proper colored emojis that will display correctly
     results_df['Predicted_Risk'] = ["🟢 LOW" if p < 0.3 else "🟡 MODERATE" if p < 0.7 else "🔴 HIGH" for p in toxicity_prob]
     return results_df, "✅ Predicted toxicity using a pre-trained simulation model.\n"
     
@@ -576,15 +545,21 @@ with tab1:
         smiles_list = [s.strip() for s in smiles_input_p1.split('\n') if s.strip()]
         props_df, log_props = calculate_molecular_properties(smiles_list)
         full_log += log_props
-        analysis_df, log_lipinski = assess_drug_likeness(props_df)
+        
+        # --- MODIFIED LOGIC ---
+        # assess_drug_likeness now returns two dataframes: one for analysis/plotting, one for display
+        analysis_df, display_df, log_lipinski = assess_drug_likeness(props_df)
         full_log += log_lipinski
+        
+        # The plotting function uses the 'analysis_df' with the boolean 'Drug_Like' column
         props_plot, log_plot = plot_properties_dashboard(analysis_df)
         full_log += log_plot
         full_log += "\n--- Phase 1 Analysis Complete ---"
         st.session_state.log_p1 = full_log
         
+        # The display dataframe has the pretty emojis
         lipinski_cols = ['Molecule', 'MW', 'LogP', 'HBD', 'HBA', 'Lipinski_Violations', 'Drug_Like']
-        lipinski_subset_df = analysis_df[lipinski_cols] if not analysis_df.empty else pd.DataFrame(columns=lipinski_cols)
+        lipinski_subset_df = display_df[lipinski_cols] if not display_df.empty else pd.DataFrame(columns=lipinski_cols)
 
         st.session_state.results_p1 = {
             'protein_view_html': protein_view_html,
@@ -601,7 +576,8 @@ with tab1:
         p1_tabs = st.tabs(["Protein Information", "Molecule Analysis", "Analysis Plots"])
         with p1_tabs[0]:
             st.subheader("Protein 3D Structure (Interactive)")
-            st.components.v1.html(res1.get('protein_view_html', '<p>No data</p>'), height=600, scrolling=False)
+            if res1.get('protein_view_html'):
+                st.components.v1.html(res1['protein_view_html'], height=600, scrolling=False)
             st.subheader("FASTA Sequence Information")
             st.text_area("", res1.get('fasta_log', 'No data'), height=200, key="fasta_info_area")
         with p1_tabs[1]:
@@ -613,6 +589,8 @@ with tab1:
             st.subheader("Molecular Properties Dashboard")
             if res1.get('props_plot'):
                 st.pyplot(res1['props_plot'])
+            else:
+                st.warning("Could not generate plots. Please check the logs for more details.")
 
 # ===== TAB 2: LEAD GENERATION & OPTIMIZATION =====
 with tab2:
@@ -668,10 +646,12 @@ with tab2:
                 st.dataframe(res2.get('admet_df', pd.DataFrame()), use_container_width=True)
         with p2_tabs[1]:
             st.subheader("Interactive 2D and 3D views of candidate molecules.")
-            st.components.v1.html(res2.get('combined_viz_html', '<p>No data</p>'), height=700, scrolling=True)
+            if res2.get('combined_viz_html'):
+                st.components.v1.html(res2.get('combined_viz_html'), height=700, scrolling=True)
         with p2_tabs[2]:
             st.subheader("Detailed view of the top candidate binding to the protein.")
-            st.components.v1.html(res2.get('interaction_html', '<p>No data</p>'), height=700, scrolling=False)
+            if res2.get('interaction_html'):
+                st.components.v1.html(res2.get('interaction_html'), height=700, scrolling=False)
 
 
 # ===== TAB 3: PRECLINICAL DEVELOPMENT =====
@@ -717,7 +697,8 @@ with tab3:
                 st.dataframe(res3.get('tox_df', pd.DataFrame()), use_container_width=True)
         with p3_tabs[1]:
             st.subheader("Interactive 3D gallery of the compounds under analysis.")
-            st.components.v1.html(res3.get('combined_viz_html', '<p>No data</p>'), height=1000, scrolling=True)
+            if res3.get('combined_viz_html'):
+                st.components.v1.html(res3.get('combined_viz_html'), height=1000, scrolling=True)
 
 
 # ===== TAB 4: POST-MARKET SURVEILLANCE =====
@@ -756,7 +737,8 @@ with tab4:
         p4_tabs = st.tabs(["Pharmacovigilance Analysis", "Regulatory & Ethical Frameworks"])
         with p4_tabs[0]:
             st.subheader("Simulated Adverse Event Analysis")
-            st.pyplot(res4['plot_bar'])
+            if res4.get('plot_bar'):
+                st.pyplot(res4['plot_bar'])
             st.dataframe(res4['rwd_df'], use_container_width=True)
             
         with p4_tabs[1]:
@@ -767,3 +749,4 @@ with tab4:
             with col2:
                 st.subheader("Ethical Framework for AI in Healthcare")
                 st.dataframe(res4.get('eth_df', pd.DataFrame()), use_container_width=True)
+