# Import necessary libraries
import streamlit as st
from io import StringIO, BytesIO
from Bio import SeqIO
from kitikiplot.genomics import grid

# Page configuration
st.set_page_config(
    page_title="Genome Grid Plot",
    page_icon="🧬",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Professional styling with blue navigation bar
st.markdown("""
<style>
    /* Remove default Streamlit padding and margins */
    .block-container {
        padding-top: 0rem;
        padding-bottom: 0rem;
        padding-left: 1rem;
        padding-right: 1rem;
    }
    
    /* Hide Streamlit header and footer */
    header[data-testid="stHeader"] {
        display: none;
    }
    
    .main > div {
        padding-top: 0rem;
    }
    
    /* Professional navigation bar */
    .nav-bar {
        background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
        padding: 0.05rem 2rem;
        color: white;
        display: flex;
        align-items: center;
        margin: -1rem -1rem 2rem -1rem;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
    }
    
    .nav-logo {
        height: 40px;
        margin-right: 15px;
        border: 1px solid #9ef2bd;
        border-radius: 2px; /* Optional: smooth corners */
    }
    
    .nav-title {
        font-size: 28px;
        font-weight: 600;
        color: white;
        margin: 0;
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    
    /* Professional card styling */
    .upload-card {
        background: white;
        border-radius: 12px;
        padding: 2rem;
        box-shadow: 0 4px 20px rgba(0,0,0,0.08);
        border: 1px solid #e8eaed;
    }
    
    .plot-card {
        background: white;
        border-radius: 12px;
        padding: 2rem;
        box-shadow: 0 4px 20px rgba(0,0,0,0.08);
        border: 1px solid #e8eaed;
        min-height: 600px;
    }
    
    /* Professional button styling */
    .stButton > button {
        background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
        color: white;
        border: none;
        border-radius: 8px;
        padding: 0.75rem 2rem;
        font-weight: 600;
        font-size: 16px;
        transition: all 0.3s ease;
        width: 100%;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(30,60,114,0.3);
    }
    
    /* Download button styling */
    .stDownloadButton > button {
        background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
        color: white;
        border: none;
        border-radius: 8px;
        padding: 0.75rem 2rem;
        font-weight: 600;
        font-size: 16px;
        transition: all 0.3s ease;
        width: 30%;
    }
    
    .stDownloadButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(40,167,69,0.3);
    }
    
    /* File uploader styling */
    .stFileUploader > div > div {
        border: 2px dashed #2a5298;
        border-radius: 8px;
        padding: 2rem;
        text-align: center;
        background: #f8f9ff;
    }
    
    /* Number input styling */
    .stNumberInput > div > div > input {
        border-radius: 8px;
        border: 2px solid #e8eaed;
        padding: 0.75rem;
        font-size: 16px;
    }
    
    /* Section headers */
    .section-header {
        color: #1e3c72;
        font-size: 24px;
        font-weight: 600;
        margin-bottom: 1rem;
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    
    /* Info text */
    .info-text {
        color: #6c757d;
        font-size: 14px;
        margin-top: 0.5rem;
    }
</style>
""", unsafe_allow_html=True)

import base64
with open("src/logo.png", "rb") as img_file:
    encoded = base64.b64encode(img_file.read()).decode()

# Display using HTML
st.markdown(f"""
<div class="nav-bar">
    <img src="data:image/png;base64,{encoded}" class="nav-logo" alt="Logo" style="height:40px;">
    <h2 class="nav-title">Genome Grid Plot</h2>
</div>
""", unsafe_allow_html=True)

# st.markdown("""
# <div class="nav-bar">
#     <h1 class="nav-title">Genome Grid Plot</h1>
# </div>
# """, unsafe_allow_html=True)

# st.image("banner.png", width=60)  

# Initialize session state
if 'figure' not in st.session_state:
    st.session_state.figure = None
if 'data' not in st.session_state:
    st.session_state.data = ""
if 'plot_generated' not in st.session_state:
    st.session_state.plot_generated = False

# Main layout
col1, col2 = st.columns([30, 70])

# Column 1: Upload and Controls
with col1:
    st.markdown('<h3 class="section-header">📁 Upload Genome File</h3>', unsafe_allow_html=True)
    
    uploaded_file = st.file_uploader(
        "Choose a FASTA file", 
        type=['fasta', 'fa', 'fas'],
        help="Upload your genome sequence in FASTA format"
    )
    
    if uploaded_file is not None:
        # Process file
        try:
            data = str(list(SeqIO.parse(StringIO(uploaded_file.getvalue().decode("utf-8")), "fasta"))[0].seq)
            st.session_state.data = data
            
            # Show file info
            st.success(f"✅ Sequence length: {len(data):,} nucleotides")
            
            # Window length input
            st.markdown('<h4 class="section-header">⚙️ Parameters</h4>', unsafe_allow_html=True)
            window_length = st.number_input(
                label="Window/Chunk Length", 
                min_value=1, 
                max_value=200, 
                value=50,
                key="window_length",
                help="Size of each analysis window for the grid plot"
            )
            st.markdown('<p class="info-text">Recommended: sequence length <1000 for detailed analysis</p>', unsafe_allow_html=True)
            
            # Submit button
            submit = st.button("🔬 Generate Plot", type="primary", key="submit_btn")
            
            if submit:
                with st.spinner("🧬 Generating genome grid plot..."):
                    st.session_state.figure = grid.plot(
                        nucleotide_sequence=data, 
                        window_length=window_length
                    )
                    st.session_state.plot_generated = True
                
        except Exception as e:
            st.error(f"❌ Error processing file: {str(e)}")
    
    else:
        st.info("👆 Please upload a FASTA file to begin visualization")
    
    # # Banner image at bottom with error handling
    # try:
    #     st.image("src/banner.png", use_container_width=True)
    # except:
    #     pass  # Silently ignore if banner not found
        
    st.markdown('</div>', unsafe_allow_html=True)

# Column 2: Plot Display
with col2:
    
    if st.session_state.figure is not None and st.session_state.plot_generated:

        st.markdown('<h3 class="section-header">Visualization Results</h3>', unsafe_allow_html=True)

        # Display plot
        st.pyplot(st.session_state.figure, use_container_width=True)
        
        # Download section
        st.markdown("---")
        
        # Prepare download
        img_bytes = BytesIO()
        st.session_state.figure.savefig(img_bytes, format="png", dpi=300, bbox_inches='tight', 
                                      facecolor='white', edgecolor='none')
        img_bytes.seek(0)
        
        # Download button
        filename = uploaded_file.name.replace('.fasta', '').replace('.fa', '').replace('.fas', '') if uploaded_file else "genome"
        st.download_button(
            label="📥 Download Plot as PNG",
            data=img_bytes.getvalue(),
            file_name=f"{filename}_GridPlot.png",
            mime="image/png"
        )
        
    else:
        # Empty state with instructions
        st.markdown("""
        <div style="text-align: center; padding: 4rem 2rem; color: #343a40; font-family: 'Segoe UI', sans-serif;">
            <h4 style="margin-bottom: 1rem;">🧬 Ready for Analysis <b>?</b></h4>
            <p style="font-size: 1rem;">Upload a FASTA file and click <strong>"Generate Plot"</strong> to visualize your genome sequence.</p>
            <br>
            <div style="background: #ffffff; padding: 2rem; border-radius: 10px; border: 1px solid #e0e0e0; box-shadow: 0 2px 8px rgba(0,0,0,0.05); max-width: 700px; margin: 0 auto;">
                <h5 style="color: #2a5298; font-size: 1.3rem; text-align: left;">About</h5>
                <ul style="text-align: left; padding-left: 1.2rem; margin: 0; line-height: 1.8; color: #495057; font-size: 0.95rem;">
                    <li>Visualize nucleotide composition patterns</li>
                    <li>Analyze sequence complexity and structure</li>
                    <li>Identify repetitive regions and motifs</li>
                    <li>Intuitive for genomics research and education</li>
                    <li>Powered by <b><i>kitikiplot</i></b></li>
                </ul>
            </div>
        </div>
        """, unsafe_allow_html=True)

    
    st.markdown('</div>', unsafe_allow_html=True)

# Footer
st.markdown("---")
st.markdown("""
<div style="text-align: center; color: #2a5298; padding: 0.01rem;">
    <p><strong>Genome Grid Plot</strong> • Intuitive genomics visualization tool • Built with Streamlit</p>
</div>
""", unsafe_allow_html=True)