import streamlit as st
import pandas as pd
import os

from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report

from utils.summarizer import summarize_genes
from utils.ml_model import train_model, get_top_biomarkers
from utils.preprocess import normalize_counts


st.title("🧬 ML-based Biomarker Discovery from RNA-seq")


uploaded_file = st.file_uploader(
    "Upload RNA-seq CSV file",
    type=["csv"]
)


if uploaded_file:

    df = pd.read_csv(uploaded_file)

    st.write("### Dataset Preview")
    st.dataframe(df.head())


    if "type" not in df.columns:

        st.error(
            "Dataset must contain 'type' column with labels"
        )

    else:

        # ----------------------------
        # Separate labels
        # ----------------------------

        labels = df["type"].tolist()

        counts = df.drop(columns=["type"])


        # ----------------------------
        # Normalize RNA counts
        # ----------------------------

        normalized = normalize_counts(
            counts
        )


        st.write("### Normalized Data")
        st.dataframe(normalized.head())


        # ----------------------------
        # Train ML Model
        # ----------------------------

        model, feature_importance = train_model(
            normalized,
            labels
        )


        st.write("### Top Biomarker Genes")

        top_genes = get_top_biomarkers(
            feature_importance,
            top_n=20
        )

        st.dataframe(top_genes)


        # ----------------------------
        # PCA Visualization
        # ----------------------------

        X = normalized.T

        scaler = StandardScaler()

        X_scaled = scaler.fit_transform(X)


        pca = PCA(
            n_components=2
        )

        X_pca = pca.fit_transform(
            X_scaled
        )


        pca_df = pd.DataFrame(
            X_pca,
            columns=[
                "PC1",
                "PC2"
            ]
        )

        pca_df["label"] = labels


        st.write("### PCA Plot")

        st.scatter_chart(
            pca_df,
            x="PC1",
            y="PC2",
            color="label"
        )


        # ----------------------------
        # Groq Gene Explanation
        # ----------------------------

        groq_api_key = os.environ.get(
            "GROQ_API_KEY"
        )


        if groq_api_key:

            summaries = summarize_genes(
                top_genes["Gene"].tolist(),
                groq_api_key
            )


            st.write(
                "### AI Biomarker Explanation"
            )


            for gene, summary in summaries.items():

                st.markdown(
                    f"**{gene}**: {summary}"
                )

        else:

            st.warning(
                "Groq API key not found"
            )