import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
from scipy import stats
import py3Dmol
from streamlit.components.v1 import html
import networkx as nx
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import classification_report, confusion_matrix
import shap # Keep import for potential future use or if other parts rely on it, though not used for the direct plot now
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
from scipy.integrate import odeint

# --- Page Configuration ---
st.set_page_config(layout="wide", page_title="Computational Biology Analysis")

# --- Global Styling ---
# Define the color scheme
BACKGROUND_COLOR = 'rgb(28, 28, 28)'
TEXT_COLOR = 'white'
ACCENT_COLOR = '#FFA500' # Orange for accents
PLOT_FACE_COLOR = '#1c1c1c'

# Custom CSS for font and dark theme
st.markdown(f"""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@400;700&display=swap');

    html, body, [class*="st-"], .main, .stApp {{
        background-color: {BACKGROUND_COLOR};
        color: {TEXT_COLOR};
        font-family: 'Roboto', 'Funnel Sans', sans-serif;
    }}
    .stTabs [data-baseweb="tab-list"] {{
        gap: 24px;
    }}
    .stTabs [data-baseweb="tab"] {{
        height: 50px;
        white-space: pre-wrap;
        background-color: transparent;
        border-radius: 4px 4px 0px 0px;
        gap: 1px;
        padding-top: 10px;
        padding-bottom: 10px;
    }}
    /* --- FIX STARTS HERE --- */
    .stTabs [aria-selected="true"] {{
        background-color: {PLOT_FACE_COLOR};
        border-bottom: 3px solid {ACCENT_COLOR};
    }}
    /* --- FIX ENDS HERE --- */
    h1, h2, h3, h4, h5, h6 {{
        color: {TEXT_COLOR};
    }}
    .stButton>button {{
        background-color: {ACCENT_COLOR};
        color: black;
        border: none;
        padding: 10px 20px;
        border-radius: 8px;
    }}
</style>
""", unsafe_allow_html=True)


# --- Data Generation and Models (from notebook) ---
@st.cache_data
def generate_ml_data():
    np.random.seed(42)
    n_samples = 300
    tnf_a = np.concatenate([np.random.normal(95, 10, 100), np.random.normal(155, 15, 100), np.random.normal(145, 15, 100)])
    myhc_ratio = np.concatenate([np.random.normal(1.1, 0.1, 100), np.random.normal(0.8, 0.1, 100), np.random.normal(1.2, 0.1, 100)])
    prob_dysfunction = 1 / (1 + np.exp(-(0.05 * (tnf_a - 150) - 7 * (myhc_ratio - 1.0))))
    systolic_dysfunction = np.random.binomial(1, prob_dysfunction)
    df_ml = pd.DataFrame({'TNF_alpha': tnf_a, 'MyHC_Ratio': myhc_ratio, 'Dysfunction': systolic_dysfunction})
    return df_ml

@st.cache_resource
def train_model(df_ml):
    X = df_ml[['TNF_alpha', 'MyHC_Ratio']]
    y = df_ml['Dysfunction']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
    model = GradientBoostingClassifier(n_estimators=100, learning_rate=0.1, max_depth=3, random_state=42)
    model.fit(X_train, y_train)
    return model, X_train, X_test, y_train, y_test

# --- ODE Model for Simulation ---
def myhc_model_odes(y, t, params):
    k_on, k_off, k_activate_tf, k_tf_bind_myh6, k_tf_unbind_myh6, k_tf_bind_myh7, k_tf_unbind_myh7, k_produce_myh6_base, k_produce_myh6_repressed, k_produce_myh7_base, k_produce_myh7_activated, k_degrade_myhc = params
    TNFa_free, TNFR_free, TNFa_TNFR, TF_inactive, TF_active_free, Gene_myh6_free, Gene_myh6_TF, Gene_myh7_free, Gene_myh7_TF, MyHC_alpha, MyHC_beta = y
    dydt = np.zeros(11)
    dydt[0] = -k_on * TNFa_free * TNFR_free + k_off * TNFa_TNFR
    dydt[1] = -k_on * TNFa_free * TNFR_free + k_off * TNFa_TNFR
    dydt[2] = k_on * TNFa_free * TNFR_free - k_off * TNFa_TNFR
    dydt[3] = -k_activate_tf * TNFa_TNFR * TF_inactive
    dydt[4] = (k_activate_tf * TNFa_TNFR * TF_inactive - k_tf_bind_myh6 * TF_active_free * Gene_myh6_free + k_tf_unbind_myh6 * Gene_myh6_TF - k_tf_bind_myh7 * TF_active_free * Gene_myh7_free + k_tf_unbind_myh7 * Gene_myh7_TF)
    dydt[5] = -k_tf_bind_myh6 * TF_active_free * Gene_myh6_free + k_tf_unbind_myh6 * Gene_myh6_TF
    dydt[6] = k_tf_bind_myh6 * TF_active_free * Gene_myh6_free - k_tf_unbind_myh6 * Gene_myh6_TF
    dydt[7] = -k_tf_bind_myh7 * TF_active_free * Gene_myh7_free + k_tf_unbind_myh7 * Gene_myh7_TF
    dydt[8] = k_tf_bind_myh7 * TF_active_free * Gene_myh7_free - k_tf_unbind_myh7 * Gene_myh7_TF
    dydt[9] = (k_produce_myh6_base * Gene_myh6_free + k_produce_myh6_repressed * Gene_myh6_TF - k_degrade_myhc * MyHC_alpha)
    dydt[10] = (k_produce_myh7_base * Gene_myh7_free + k_produce_myh7_activated * Gene_myh7_TF - k_degrade_myhc * MyHC_beta)
    return dydt

@st.cache_data
def run_ode_simulation():
    params = [1.0, 0.1, 1e-2, 1e-3, 0.5, 1e-3, 0.5, 1e-1, 1e-3, 1e-2, 1e-1, 5e-3]
    y0 = [100, 20, 0, 50, 0, 10, 0, 10, 0, 0, 0]
    t = np.linspace(0, 5000, 1001)
    solution = odeint(myhc_model_odes, y0, t, args=(params,))
    return t, solution

# --- Plotting Functions ---
def render_py3dmol(pdb_id, chain_styles, label_text, label_options, bg_color='0x1c1c1c'):
    view = py3Dmol.view(query=f'pdb:{pdb_id}', width='100%', height=400)
    for chain, style in chain_styles.items():
        view.setStyle({chain: chain}, style)
    view.addLabel(label_text, {'fontColor':'white', 'backgroundColor':'#333333', 'fontSize':16, **label_options.get('style', {})}, label_options.get('position', {}))
    view.setBackgroundColor(bg_color)
    view.zoomTo()
    # Use _make_html() to generate the HTML string for streamlit
    html_string = view._make_html()
    html(html_string, height=400)

# --- App Layout ---
st.title("TNF-α Mediated Cardiac Dysfunction: A Computational Analysis")
st.markdown("This application explores the findings of Manilall et al. (2023), who investigated how Tumor Necrosis Factor-alpha (TNF-α) mediates early-stage left ventricular (LV) systolic dysfunction. We replicate key findings and enhance the analysis with computational biology tools.")

tab1, tab2 = st.tabs(["General Lab / Analysis", "AI & Machine Learning"])

# ==============================================================================
#                                 GENERAL ANALYSIS TAB
# ==============================================================================
with tab1:
    st.header("Visual and Systems Biology Analysis")
    st.markdown("---")

    # --- TNF-α Visualization ---
    with st.container():
        st.subheader("3D Visualization of the Mediator: TNF-α")
        st.write("TNF-α is a cytokine crucial to the immune response. In chronic inflammatory conditions, its elevated levels can have damaging effects. It typically functions as a homotrimer (a complex of three identical protein chains). We visualize the PDB structure 1TNF.")
        
        chain_styles_tnf = {
            'A': {'cartoon': {'color': 'salmon'}},
            'B': {'cartoon': {'color': 'lightblue'}},
            'C': {'cartoon': {'color': 'lightgreen'}}
        }
        render_py3dmol('1TNF', chain_styles_tnf, "TNF-α Trimer", {'position': {'resi': 50, 'chain': 'A'}})
    st.markdown("---")

    # --- Myosin Visualization ---
    with st.container():
        st.subheader("The Molecular Target: Myosin Heavy Chain (MyHC)")
        st.write("The study's central finding is a TNF-α-induced switch from the faster α-MyHC (`Myh6`) to the slower β-MyHC (`Myh7`), impairing heart function. We visualize the human cardiac myosin motor domain (PDB: 5N6A).")
        
        # Use st.tabs for Myosin View to allow for active tab styling
        myosin_tab1, myosin_tab2 = st.tabs(["3D Model", "Functional Info"])

        with myosin_tab1:
             chain_styles_myosin = {
                'protein': {'cartoon': {'color':'#A9A9A9'}},
             }
             view = py3Dmol.view(query='pdb:5N6A', width='100%', height=450)
             view.setStyle({'cartoon': {'color':'#A9A9A9'}})
             view.addStyle({'resn': 'ADP'}, {'stick': {'colorscheme': 'magentaCarbon'}})
             view.addStyle({'resn': 'MG'}, {'sphere': {'color': 'purple'}})
             view.addLabel("Myosin Motor Domain", {'fontColor':'white', 'backgroundColor':'#333333'}, {'resi': 1, 'chain': 'A'})
             view.addLabel("ADP (Energy)", {'fontColor':'magenta'}, {'resn': 'ADP'})
             view.setBackgroundColor('0x1c1c1c')
             view.zoomTo()
             # Use _make_html() to generate the HTML string for streamlit
             html(view._make_html(), height=450)

        with myosin_tab2:
            st.info("""
            **α-MyHC (Myh6) Function:** Major contractile protein in the atria. In the ventricle, it is the predominant isoform in small mammals and in fetal human heart but it is replaced by MYH7 after birth. Contributes to cardiac muscle contraction.
            
            **β-MyHC (Myh7) Function:** Muscle contraction. The major contractile protein in the ventricular myocardium.
            """)
    st.markdown("---")

    # --- Heatmap ---
    with st.container():
        st.subheader("Integrated Heatmap of Key Study Findings")
        st.write("A heatmap consolidating data on inflammation, cardiac function, and gene expression. The anti-TNF-α treatment clearly reverses the pathological changes seen in the CIA group.")
        
        data = {
            'Circulating TNF-α (pg/mL)': [93.9, 155.8, 145.6],
            'Global Longitudinal Strain (%)': [-19.6, -16.7, -19.0],
            'Myh6 mRNA (α-MyHC)': [1.03, 0.89, 1.11],
            'Myh7 mRNA (β-MyHC)': [0.96, 1.25, 0.91],
            'Myh6/Myh7 Ratio': [1.13, 0.79, 1.24],
            'IL-6 mRNA': [0.79, 1.65, 1.15]
        }
        df_heatmap = pd.DataFrame(data, index=['Control', 'CIA', 'CIA+anti-TNF-α'])
        df_normalized = (df_heatmap - df_heatmap.mean()) / df_heatmap.std()

        fig_heatmap = px.imshow(df_normalized.T,
                                 text_auto=True,
                                 aspect="auto",
                                 color_continuous_scale='RdBu_r', # Red-Blue reversed for coolwarm effect
                                 labels=dict(x="Experimental Group", y="Parameter", color="Normalized Value"),
                                 x=df_heatmap.index,
                                 y=df_heatmap.columns)

        fig_heatmap.update_layout(
            title_text='Integrated Heatmap of Key Study Findings',
            title_x=0.5,
            height=600,
            xaxis_title='Experimental Group',
            yaxis_title='Parameter',
            paper_bgcolor=BACKGROUND_COLOR,
            plot_bgcolor=PLOT_FACE_COLOR,
            font=dict(color=TEXT_COLOR),
            coloraxis_colorbar=dict(title="Normalized Value", thicknessmode="pixels", thickness=20, lenmode="pixels", len=300)
        )
        # Add original values as text on hover
        fig_heatmap.update_traces(text=df_heatmap.T.values.round(2), texttemplate="%{text}", hovertemplate="<b>%{y}</b><br>Group: %{x}<br>Value: %{text}<extra></extra>")
        
        st.plotly_chart(fig_heatmap, use_container_width=True)
    st.markdown("---")

    # --- Pathway Simulation ---
    with st.container():
        st.subheader("Simulating the TNF-α Signaling Pathway")
        st.write("A simplified computational model to simulate the paper's central hypothesis: increased TNF-α leads to a decrease in the Myh6/Myh7 ratio over time.")
        
        t, solution = run_ode_simulation()
        MyHC_alpha = solution[:, 9]
        MyHC_beta = solution[:, 10]
        TF_active_total = solution[:, 4] + solution[:, 6] + solution[:, 8]

        # Create subplots
        fig_sim = make_subplots(rows=2, cols=1,
                                shared_xaxes=True,
                                vertical_spacing=0.1,
                                subplot_titles=("Simulated MyHC Isoform Switch Induced by TNF-α", "Transcription Factor Dynamics"))

        # MyHC Isoform Switch plot (top subplot)
        fig_sim.add_trace(go.Scatter(x=t, y=MyHC_alpha, mode='lines', name='α-MyHC (Myh6)',
                                     line=dict(color='blue', width=2)),
                          row=1, col=1)
        fig_sim.add_trace(go.Scatter(x=t, y=MyHC_beta, mode='lines', name='β-MyHC (Myh7)',
                                     line=dict(color='red', width=2, dash='dash')),
                          row=1, col=1)
        fig_sim.update_yaxes(title_text='MyHC Protein Level (Arbitrary Units)', row=1, col=1)
        fig_sim.update_xaxes(title_text='Time (Arbitrary Units)', row=1, col=1, showticklabels=False) # Hide x-axis labels on top plot

        # Transcription Factor Dynamics subplot (bottom subplot)
        fig_sim.add_trace(go.Scatter(x=t, y=solution[:, 2], mode='lines', name='TNF-α:TNFR Complex',
                                     line=dict(color='green', width=2)),
                          row=2, col=1)
        fig_sim.add_trace(go.Scatter(x=t, y=TF_active_total, mode='lines', name='Active TF (Total)',
                                     line=dict(color='orange', width=2)),
                          row=2, col=1)
        fig_sim.update_yaxes(title_text='Concentration (Arbitrary Units)', row=2, col=1)
        fig_sim.update_xaxes(title_text='Time (Arbitrary Units)', row=2, col=1)

        fig_sim.update_layout(
            height=700, # Increased height to accommodate two distinct plots
            paper_bgcolor=BACKGROUND_COLOR,
            plot_bgcolor=PLOT_FACE_COLOR,
            font=dict(color=TEXT_COLOR),
            legend=dict(x=1.02, y=1, yanchor="top", xanchor="left", bgcolor='rgba(0,0,0,0)'), # Consolidated legend
            margin=dict(l=40, r=40, t=80, b=40)
        )
        
        st.plotly_chart(fig_sim, use_container_width=True)
    st.markdown("---")

    # --- Network Analysis ---
    with st.container():
        st.subheader("Network Analysis of Pathological Relationships")
        st.write("A high-level overview of the proposed pathophysiology, from inflammation to systolic dysfunction.")
        
        G = nx.DiGraph()
        nodes_data = [
            ("Systemic Inflammation\n(CIA Model)", 'red', 3500, 0.0), ("↑ Circulating TNF-α", 'orange', 3500, 0.5),
            ("↓ Myh6/Myh7 Ratio", 'purple', 3500, 0.5), ("Impaired Systolic Function\n(↓ Strain & Velocity)", 'darkred', 4000, -0.5),
            ("↑ IL-6 Expression", 'gold', 3000, 0.5), ("Anti-TNF-α\nTreatment", 'green', 3000, 1.0)
        ]
        for name, color, size, z in nodes_data:
            G.add_node(name, color=color, size=size, z=z)
        G.add_edge("Systemic Inflammation\n(CIA Model)", "↑ Circulating TNF-α")
        G.add_edge("↑ Circulating TNF-α", "↓ Myh6/Myh7 Ratio", label="mediates")
        G.add_edge("↑ Circulating TNF-α", "↑ IL-6 Expression", label="mediates")
        G.add_edge("↓ Myh6/Myh7 Ratio", "Impaired Systolic Function\n(↓ Strain & Velocity)", label="underlies")
        G.add_edge("Anti-TNF-α\nTreatment", "↓ Myh6/Myh7 Ratio", label="prevents", type='inhibit')
        G.add_edge("Anti-TNF-α\nTreatment", "Impaired Systolic Function\n(↓ Strain & Velocity)", label="prevents", type='inhibit')

        
        pos_2d = nx.spring_layout(G, k=1.8, seed=42)
        node_trace = go.Scatter3d(
            x=[pos_2d[node][0] for node in G.nodes()], y=[pos_2d[node][1] for node in G.nodes()], z=[G.nodes[node]['z'] for node in G.nodes()],
            mode='markers+text',
            marker=dict(size=[G.nodes[node]['size']/300 for node in G.nodes()], color=[G.nodes[node]['color'] for node in G.nodes()], opacity=0.9),
            text=list(G.nodes()), textposition="top center", textfont=dict(size=10, color=TEXT_COLOR), name="Nodes"
        )
        edge_traces = []
        for edge in G.edges():
            x0, y0 = pos_2d[edge[0]]; x1, y1 = pos_2d[edge[1]]
            z0, z1 = G.nodes[edge[0]]['z'], G.nodes[edge[1]]['z']
            color = 'green' if G.edges[edge].get('type') == 'inhibit' else TEXT_COLOR
            dash = 'dash' if G.edges[edge].get('type') == 'inhibit' else 'solid'
            edge_trace = go.Scatter3d(x=[x0, x1, None], y=[y0, y1, None], z=[z0, z1, None], mode='lines', line=dict(color=color, width=4, dash=dash), showlegend=False)
            edge_traces.append(edge_trace)

        fig_3d = go.Figure(data=[node_trace] + edge_traces)
        fig_3d.update_layout(
            title="Interactive 3D Network",
            height=600, 
            scene=dict(xaxis_title="X", yaxis_title="Y", zaxis_title="Hierarchy",
                       xaxis=dict(showbackground=False, visible=False),
                       yaxis=dict(showbackground=False, visible=False),
                       zaxis=dict(showbackground=False, visible=False)),
            paper_bgcolor=BACKGROUND_COLOR,
            plot_bgcolor=BACKGROUND_COLOR,
            font=dict(color=TEXT_COLOR),
            showlegend=False,
            margin=dict(l=0, r=0, b=0, t=40)
        )
        
        st.plotly_chart(fig_3d, use_container_width=True)


# ==============================================================================
#                               AI & ML TAB
# ==============================================================================
with tab2:
    st.header("AI & Machine Learning Insights")
    st.markdown("---")

    df_ml = generate_ml_data()
    model, X_train, X_test, y_train, y_test = train_model(df_ml)

    # --- Predictive Modeling ---
    with st.container():
        st.subheader("Predictive Modeling: Can Biomarkers Predict Dysfunction?")
        st.write("We train a Gradient Boosting model to predict systolic dysfunction based on TNF-α levels and the Myh6/Myh7 ratio. Below is the model's performance on the test dataset.")
        
        y_pred = model.predict(X_test)
        
        col1, col2 = st.columns(2)
        with col1:
            st.text("Classification Report")
            report = classification_report(y_test, y_pred, target_names=['Normal', 'Dysfunction'], output_dict=True)
            st.dataframe(pd.DataFrame(report).transpose())
        
        with col2:
            cm = confusion_matrix(y_test, y_pred, labels=model.classes_)
            fig_cm = px.imshow(cm,
                                text_auto=True,
                                labels=dict(x="Predicted", y="True", color="Count"),
                                x=['Normal', 'Dysfunction'],
                                y=['Normal', 'Dysfunction'],
                                color_continuous_scale='Blues')

            fig_cm.update_layout(
                title_text='Confusion Matrix',
                title_x=0.5,
                xaxis_title="Predicted Label",
                yaxis_title="True Label",
                paper_bgcolor=BACKGROUND_COLOR,
                plot_bgcolor=PLOT_FACE_COLOR,
                font=dict(color=TEXT_COLOR)
            )
            st.plotly_chart(fig_cm, use_container_width=True)
    st.markdown("---")

    # --- Explainable AI (XAI) ---
    with st.container():
        st.subheader("Traditional Feature Importance")
        st.write("Here, we visualize the importance of each feature as determined by the Gradient Boosting model. Features with higher importance contribute more significantly to the model's predictions.")
        
        # Get feature importances
        feature_importances = pd.DataFrame({
            'Feature': X_train.columns,
            'Importance': model.feature_importances_
        }).sort_values('Importance', ascending=False)

        # Create bar chart for feature importance
        fig_importance = px.bar(feature_importances,
                                x='Importance',
                                y='Feature',
                                orientation='h',
                                title='Feature Importance from Gradient Boosting Model',
                                color_discrete_sequence=[ACCENT_COLOR]) # Use ACCENT_COLOR for bars
        
        fig_importance.update_layout(
            paper_bgcolor=BACKGROUND_COLOR,
            plot_bgcolor=PLOT_FACE_COLOR,
            font=dict(color=TEXT_COLOR),
            xaxis_title='Importance',
            yaxis_title='Feature',
            yaxis={'categoryorder':'total ascending'} # Ensure the most important features are at the top
        )
        
        st.plotly_chart(fig_importance, use_container_width=True)

    st.markdown("---")
    # --- Clustering ---
    with st.container():
        st.subheader("Unsupervised Learning: Identifying Patient Subgroups")
        st.write("K-Means clustering can identify natural groupings in the data without predefined labels. This technique could help discover patient subgroups in a real clinical dataset. The clusters found here align well with the original experimental groups (Control, CIA, Treated).")
        
        X_cluster = df_ml[['TNF_alpha', 'MyHC_Ratio']]
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X_cluster)
        kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
        df_ml['Cluster'] = kmeans.fit_predict(X_scaled).astype(str) # Convert to string for discrete colors

        centroids = scaler.inverse_transform(kmeans.cluster_centers_)
        centroids_df = pd.DataFrame(centroids, columns=['TNF_alpha', 'MyHC_Ratio'])
        centroids_df['Cluster'] = ['Centroid 0', 'Centroid 1', 'Centroid 2'] # Label centroids

        fig_cluster = px.scatter(
            df_ml, x='TNF_alpha', y='MyHC_Ratio', color='Cluster',
            color_discrete_sequence=px.colors.qualitative.Vivid, # Viridis like qualitative palette
            title='K-Means Clustering of Biomarker Data',
            labels={'TNF_alpha': 'Circulating TNF-α (pg/mL)', 'MyHC_Ratio': 'Myh6/Myh7 Ratio'}
        )
        
        fig_cluster.add_trace(
            go.Scatter(
                x=centroids_df['TNF_alpha'],
                y=centroids_df['MyHC_Ratio'],
                mode='markers',
                marker=dict(symbol='x', size=15, color='red', line=dict(width=2, color='darkred')),
                name='Centroids',
                showlegend=True
            )
        )

        # Make the plot square
        fig_cluster.update_layout(
            paper_bgcolor=BACKGROUND_COLOR,
            plot_bgcolor=PLOT_FACE_COLOR,
            font=dict(color=TEXT_COLOR),
            xaxis_title='Circulating TNF-α (pg/mL)',
            yaxis_title='Myh6/Myh7 Ratio',
            height=600,  # Set height to be equal to width for a square appearance
            width=600,   # Set width explicitly for square appearance
            showlegend=True
        )
        st.plotly_chart(fig_cluster, use_container_width=True)