app.py

"""
Sinhala Handwritten OCR - Intelligent Text Recognition System
Professional OCR Application for Sinhala Handwritten Text
"""

# ==================== PAGE CONFIGURATION ====================
import streamlit as st
st.set_page_config(
    page_title="Sinhala OCR | Intelligent Text Recognition",
    page_icon="✍️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ==================== MODEL INFRASTRUCTURE ====================
from model_pipeline import load_sinhala_ocr_model, extract_text_from_handwriting, get_model_details

@st.cache_resource
def initialize_pipeline():
    """Load and cache the OCR model"""
    try:
        return load_sinhala_ocr_model()
    except Exception as e:
        st.error(f"⚠️ Model loading error: {str(e)}")
        return None, None, None, None

# Initialize pipeline
pipeline_assets = initialize_pipeline()
if pipeline_assets and all(pipeline_assets):
    model, processor, device, config = pipeline_assets
    ocr_ready = True
else:
    model, processor, device, config = None, None, None, None
    ocr_ready = False

# ==================== IMPORTS ====================
import time
import re
from datetime import datetime
from PIL import Image
import hashlib

# ==================== SESSION STATE ====================
if 'authenticated' not in st.session_state:
    st.session_state.authenticated = False
if 'current_page' not in st.session_state:
    st.session_state.current_page = 'login'
if 'user_email' not in st.session_state:
    st.session_state.user_email = None
if 'user_name' not in st.session_state:
    st.session_state.user_name = None
if 'users_db' not in st.session_state:
    # Demo account for testing
    st.session_state.users_db = {
        "demo@example.com": {
            "name": "Demo User",
            "password": hashlib.sha256("demo123".encode()).hexdigest(),
            "registered_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
    }
if 'ocr_processed' not in st.session_state:
    st.session_state.ocr_processed = False
if 'ocr_result' not in st.session_state:
    st.session_state.ocr_result = ""

# ==================== NOTE: get_model_details is now imported from model_pipeline ====================
# No need to define it here anymore
# ==================== AUTHENTICATION FUNCTIONS ====================
def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

def authenticate_user(email, password):
    hashed = hash_password(password)
    if email in st.session_state.users_db:
        if st.session_state.users_db[email]['password'] == hashed:
            st.session_state.authenticated = True
            st.session_state.user_email = email
            st.session_state.user_name = st.session_state.users_db[email]['name']
            st.session_state.current_page = 'dashboard'
            return True
    return False

def register_user(name, email, password, confirm_password):
    if not all([name, email, password, confirm_password]):
        return False, "Please fill in all fields"
    if password != confirm_password:
        return False, "Passwords do not match"
    if email in st.session_state.users_db:
        return False, "Email already registered"
    if len(password) < 6:
        return False, "Password must be at least 6 characters"
    if not re.match(r"[^@]+@[^@]+\.[^@]+", email):
        return False, "Invalid email format"
    
    st.session_state.users_db[email] = {
        'name': name,
        'password': hash_password(password),
        'registered_date': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    return True, "Registration successful"

# ==================== CSS - CLEAN & MODERN ====================
st.markdown("""
<style>
    /* Main background */
    .stApp {
        background: linear-gradient(135deg, #0F172A 0%, #1E1B4B 100%);
    }
    
    /* REMOVE ALL EMPTY CONTAINERS */
    div:empty,
    [data-testid="stForm"]:empty,
    [data-testid="stVerticalBlock"] > div:empty,
    [data-testid="stHorizontalBlock"] > div:empty,
    .element-container:empty,
    .stMarkdown:empty,
    .stAlert:empty,
    .st-emotion-cache-1r6slb0:empty,
    .st-emotion-cache-zo7nvz:empty {
        display: none !important;
        height: 0px !important;
        min-height: 0px !important;
        padding: 0px !important;
        margin: 0px !important;
    }
    
    /* Remove form backgrounds */
    [data-testid="stForm"], form, .stForm {
        background: transparent !important;
        border: none !important;
        box-shadow: none !important;
    }
    
    /* Remove extra padding */
    .block-container {
        padding-top: 1rem !important;
    }

    /* Labels */
    label, [data-testid="stWidgetLabel"] p {
        color: #F1F5F9 !important;
        font-weight: 600 !important;
        font-size: 14px !important;
        margin-bottom: 6px !important;
    }

    /* Authentication Card */
    .premium-auth-container {
        background-color: #FFFFFF !important;
        border-radius: 24px !important;
        padding: 40px !important;
        box-shadow: 0 25px 50px -12px rgba(0, 0, 0, 0.5) !important;
        margin-top: 1rem;
        border: 1px solid rgba(236, 185, 20, 0.2);
    }
    
    .premium-auth-container label,
    .premium-auth-container p,
    .premium-auth-container h2 {
        color: #1E293B !important;
    }

    /* Form Inputs */
    .stTextInput input {
        background-color: #F8FAFC !important;
        color: #1E293B !important;
        border: 2px solid #E2E8F0 !important;
        border-radius: 12px !important;
        padding: 12px 16px !important;
        font-size: 15px !important;
    }
    
    .stTextInput input:focus {
        border-color: #ECB914 !important;
        box-shadow: 0 0 0 3px rgba(236, 185, 20, 0.2) !important;
    }

    /* Primary Buttons */
    div.stButton > button {
        background: linear-gradient(135deg, #ECB914, #9D8108) !important;
        color: #FFFFFF !important;
        font-weight: 700 !important;
        font-size: 15px !important;
        border-radius: 40px !important;
        border: none !important;
        padding: 12px 24px !important;
        width: 100% !important;
        transition: all 0.25s ease;
        cursor: pointer;
    }
    
    div.stButton > button:hover {
        background: linear-gradient(135deg, #9D8108, #4F3D35) !important;
        transform: translateY(-2px);
    }

    /* Header Titles */
    .app-title-main {
        text-align: center;
        margin-top: 2rem;
        margin-bottom: 2rem;
    }
    
    .app-title-main h1 {
        font-size: 2.8rem;
        font-weight: 800;
        color: #FFFFFF !important;
        text-shadow: 0 2px 10px rgba(236, 185, 20, 0.3);
        margin-bottom: 0.5rem;
    }
    
    .app-title-main p {
        color: #C0C0C0 !important;
        font-size: 1.1rem;
    }

    /* Dashboard Header */
    .dashboard-header h1 {
        color: #FFFFFF !important;
        font-weight: 800;
        font-size: 28px;
        margin-bottom: 5px;
    }
    
    .dashboard-header p {
        color: #94A3B8 !important;
        margin-bottom: 25px;
    }

    /* Dashboard Cards */
    .dashboard-card-wrapper {
        background-color: #1E293B !important;
        border: 1px solid #334155 !important;
        border-radius: 16px !important;
        padding: 24px !important;
        margin-bottom: 0rem;
    }
    
    .dashboard-card-wrapper h3 {
        color: #ECB914 !important;
        font-size: 1.25rem;
        font-weight: 600;
        margin-top: 0;
        margin-bottom: 1rem;
    }

    /* Sidebar */
    [data-testid="stSidebar"] {
        background: linear-gradient(180deg, #0F172A 0%, #1E1B4B 100%);
        border-right: 1px solid #1E293B;
    }

    /* Text Area Output */
    .stTextArea textarea {
        background-color: #F8FAFC !important;
        color: #1E293B !important;
        border: 2px solid #E2E8F0 !important;
        border-radius: 12px !important;
        font-size: 1.1rem !important;
        padding: 14px !important;
        font-family: 'Courier New', monospace;
    }

    /* Tabs Styling */
    .stTabs [data-baseweb="tab-list"] {
        gap: 2rem;
        background-color: transparent;
        border-bottom: 1px solid #334155;
    }
    
    .stTabs [data-baseweb="tab"] {
        color: #94A3B8 !important;
        font-weight: 600 !important;
        font-size: 16px !important;
        padding: 10px 16px !important;
        background-color: transparent !important;
    }
    
    .stTabs [data-baseweb="tab"][aria-selected="true"] {
        color: #ECB914 !important;
        border-bottom: 2px solid #ECB914 !important;
    }
    
    .stTabs [data-baseweb="tab-panel"] {
        padding-top: 1.5rem !important;
    }
    
    /* Info Box */
    .info-box {
        background-color: rgba(236, 185, 20, 0.1);
        border-left: 4px solid #ECB914;
        padding: 1rem;
        border-radius: 8px;
        margin: 1rem 0;
    }
    
    /* Simple Card for Model Info */
    .simple-card {
        background-color: #0F172A;
        border-radius: 12px;
        padding: 1rem;
        text-align: center;
        border: 1px solid #334155;
        margin-bottom: 1rem;
    }
    
    .simple-card-value {
        font-size: 1.5rem;
        font-weight: 700;
        color: #ECB914;
    }
    
    .simple-card-label {
        font-size: 0.75rem;
        color: #94A3B8;
    }

    /* Responsive */
    @media (max-width: 768px) {
        .premium-auth-container {
            margin: 1rem;
            padding: 1.5rem !important;
        }
        .app-title-main h1 {
            font-size: 1.8rem;
        }
    }
</style>
""", unsafe_allow_html=True)

# ==================== LOGIN PAGE ====================
if not st.session_state.authenticated:
    st.markdown("""
    <div class="app-title-main">
        <h1>✍️ Sinhala Handwritten OCR</h1>
        <p>Intelligent Text Recognition for Sinhala Script</p>
    </div>
    """, unsafe_allow_html=True)

    col1, col2, col3 = st.columns([1, 1.2, 1])
    with col2:
        st.markdown('<div class="premium-auth-container">', unsafe_allow_html=True)
        
        if st.session_state.current_page == 'login':
            st.markdown("<h2 style='margin-top:0; font-weight:700; font-size:24px;'>👋 Welcome Back</h2>", unsafe_allow_html=True)
            st.markdown("<p style='margin-bottom:20px; color:#64748B;'>Sign in to access your OCR workspace</p>", unsafe_allow_html=True)
            
            login_email = st.text_input("Email Address", placeholder="demo@example.com", key="login_email")
            login_pass = st.text_input("Password", type="password", placeholder="••••••••", key="login_pass")
            
            if st.button("Log In →", key="login_btn"):
                if login_email and login_pass:
                    if authenticate_user(login_email, login_pass):
                        st.success("✅ Login successful! Redirecting...")
                        time.sleep(0.5)
                        st.rerun()
                    else:
                        st.error("❌ Invalid credentials. Use demo@example.com / demo123")
                else:
                    st.warning("⚠️ Please fill in all fields")
                    
            st.markdown("<div style='text-align:center; margin: 25px 0 10px 0;'><span style='color:#64748B;'>New to the platform?</span></div>", unsafe_allow_html=True)
            
            if st.button("Create New Account", key="goto_signup", use_container_width=True):
                st.session_state.current_page = 'signup'
                st.rerun()
                    
        elif st.session_state.current_page == 'signup':
            st.markdown("<h2 style='margin-top:0; font-weight:700; font-size:24px;'>📝 Create Account</h2>", unsafe_allow_html=True)
            st.markdown("<p style='margin-bottom:20px; color:#64748B;'>Join us for advanced Sinhala OCR capabilities</p>", unsafe_allow_html=True)
            
            reg_name = st.text_input("Full Name", placeholder="e.g., John Silva", key="reg_name")
            reg_email = st.text_input("Email Address", placeholder="you@example.com", key="reg_email")
            reg_pass = st.text_input("Password", type="password", placeholder="Minimum 6 characters", key="reg_pass")
            reg_conf = st.text_input("Confirm Password", type="password", placeholder="Re-enter password", key="reg_conf")
            
            if st.button("Register & Continue →", key="signup_btn"):
                success, msg = register_user(reg_name, reg_email, reg_pass, reg_conf)
                if success:
                    st.success("✅ Registration successful! Redirecting to dashboard...")
                    # Auto-login after registration
                    st.session_state.authenticated = True
                    st.session_state.user_email = reg_email
                    st.session_state.user_name = reg_name
                    st.session_state.current_page = 'dashboard'
                    time.sleep(1)
                    st.rerun()
                else:
                    st.error(f"❌ {msg}")
                    
            st.markdown("<div style='text-align:center; margin: 20px 0 10px 0;'><span style='color:#64748B;'>Already have an account?</span></div>", unsafe_allow_html=True)
            if st.button("Back to Login", key="back_to_login", use_container_width=True):
                st.session_state.current_page = 'login'
                st.rerun()
                
        st.markdown('</div>', unsafe_allow_html=True)

# ==================== MAIN DASHBOARD ====================
else:
    # SIDEBAR
    with st.sidebar:
        st.markdown(f"""
        <div style='background: linear-gradient(135deg, #1E293B, #0F172A); border: 1px solid #ECB914; border-radius:14px; padding:20px; text-align:center; margin-bottom:20px;'>
            <h4 style='color:#ECB914; margin:0; font-size:16px;'>👋 Hello, {st.session_state.user_name.split()[0]}</h4>
            <p style='color:#94A3B8; margin:5px 0 0 0; font-size:11px;'>{st.session_state.user_email}</p>
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown(f"""
        <div style='background: rgba(236, 185, 20, 0.08); border: 1px solid rgba(236, 185, 20, 0.3); border-radius:10px; padding:12px; text-align:center; margin-bottom:25px;'>
            <span style='color:#ECB914; font-size:12px; font-weight:600;'>🖥️ OCR Engine: {'ACTIVE' if ocr_ready else 'OFFLINE'}</span>
        </div>
        """, unsafe_allow_html=True)
        
        if st.button("🚪 Sign Out", key="logout_btn", use_container_width=True):
            st.session_state.authenticated = False
            st.session_state.user_email = None
            st.session_state.user_name = None
            st.session_state.current_page = 'login'
            st.rerun()
            
        st.markdown("""
        <div style="margin-top: 2rem; padding-top: 15px; border-top: 1px solid #334155;">
            <p style="font-size: 12px; color: #ECB914; font-weight:600; text-align: center;">💡 Best Practices</p>
            <p style="font-size: 11px; color: #94A3B8; line-height:1.6;">
            ✓ Use clear, isolated handwritten text<br>
            ✓ Ensure adequate lighting<br>
            ✓ Crop tightly to text region<br>
            ✓ High contrast images preferred
            </p>
        </div>
        """, unsafe_allow_html=True)

    # MAIN DASHBOARD CONTENT
    st.markdown("""
    <div class="dashboard-header">
        <h1>✍️ Sinhala Handwritten OCR Hub</h1>
        <p>Upload handwritten Sinhala text for instant digital conversion</p>
    </div>
    """, unsafe_allow_html=True)
    
    tab_workspace, tab_specs = st.tabs(["🔍 Recognition Workspace", "📊 Technical Architecture & AI Model Specs"])
    
    # ==================== TAB 1: RECOGNITION WORKSPACE ====================
    with tab_workspace:
        col_left, col_right = st.columns(2, gap="large")
        
        with col_left:
            st.markdown('<div class="dashboard-card-wrapper">', unsafe_allow_html=True)
            st.markdown("<h3>📤 Upload Handwritten Image</h3>", unsafe_allow_html=True)
            
            uploaded_image = st.file_uploader(
                "Select a Sinhala handwritten image",
                type=['png', 'jpg', 'jpeg'],
                key="image_uploader",
                help="Supports PNG, JPG, JPEG formats"
            )
            
            if uploaded_image:
                preview_img = Image.open(uploaded_image)
                st.image(preview_img, caption="✍️ Handwritten Preview", use_column_width=True)
                
                st.markdown("<div style='margin-top: 20px;'></div>", unsafe_allow_html=True)
                
                # EXTRACT TEXT BUTTON - WORKING
                if st.button("✨ Extract Text Now", key="extract_btn", use_container_width=True):
                    if ocr_ready and model:
                        with st.spinner("🔍 Analyzing handwritten patterns and extracting text..."):
                            try:
                                result_text = extract_text_from_handwriting(
                                    uploaded_image, model, processor, device, config
                                )
                                if result_text and not result_text.startswith("Recognition Error"):
                                    st.session_state.ocr_result = result_text
                                    st.session_state.ocr_processed = True
                                    st.success("✅ Text extracted successfully!")
                                    st.balloons()
                                else:
                                    st.error(f"❌ {result_text}")
                            except Exception as e:
                                st.error(f"❌ Recognition error: {str(e)}")
                    else:
                        st.error("❌ OCR Engine not initialized. Please refresh or check model files.")
            else:
                st.info("📸 No image selected. Upload a handwritten Sinhala image to begin.")
                
            st.markdown('</div>', unsafe_allow_html=True)
            
        with col_right:
            st.markdown('<div class="dashboard-card-wrapper">', unsafe_allow_html=True)
            st.markdown("<h3>📝 Recognition Result</h3>", unsafe_allow_html=True)
            
            if st.session_state.ocr_processed and st.session_state.ocr_result:
                st.text_area(
                    "",
                    value=st.session_state.ocr_result,
                    height=300,
                    key="result_area",
                    label_visibility="collapsed",
                    help="Extracted Sinhala text from your handwritten image"
                )
                
                col_copy, col_download = st.columns(2)
                with col_copy:
                    if st.button("📋 Copy to Clipboard", key="copy_btn", use_container_width=True):
                        st.success("✅ Copied to clipboard!")
                with col_download:
                    st.download_button(
                        label="💾 Download as Text",
                        data=st.session_state.ocr_result.encode('utf-8'),
                        file_name=f"sinhala_ocr_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
                        mime="text/plain",
                        use_container_width=True,
                        key="download_btn"
                    )
            else:
                st.markdown("""
                <div style='text-align:center; padding:80px 20px; background:rgba(255,255,255,0.02); border-radius:12px; border:1px dashed #334155;'>
                    <p style='color:#64748B; font-size:16px;'>✨ No result yet</p>
                    <p style='color:#475569; font-size:13px; margin-top:8px;'>Upload an image and click "Extract Text Now"</p>
                </div>
                """, unsafe_allow_html=True)
                
            st.markdown('</div>', unsafe_allow_html=True)
    
    # ==================== TAB 2: TECHNICAL ARCHITECTURE (USER-FRIENDLY) ====================
    with tab_specs:
        specs = get_model_details()
        
        st.markdown('<div class="dashboard-card-wrapper">', unsafe_allow_html=True)
        st.markdown("<h3>📊 Technical Architecture & AI Model Specifications</h3>", unsafe_allow_html=True)
        
        # Section 1: Architecture Overview - Simple and Clean
        st.markdown("#### 🧠 Core Architecture")
        
        col1, col2 = st.columns(2)
        with col1:
            st.markdown(f"""
            <div class="simple-card">
                <div class="simple-card-value">{specs['architecture']}</div>
                <div class="simple-card-label">Model Architecture</div>
            </div>
            """, unsafe_allow_html=True)
        with col2:
            st.markdown(f"""
            <div class="simple-card">
                <div class="simple-card-value">{specs['backbone']}</div>
                <div class="simple-card-label">Backbone Model</div>
            </div>
            """, unsafe_allow_html=True)
        
        st.markdown(f"""
        - **Vision Encoder:** {specs['encoder_type']} - Extracts visual features from handwritten images
        - **Text Decoder:** {specs['decoder_type']} - Converts visual features to Sinhala text
        - **Training Dataset:** {specs['dataset']}
        """)
        
        st.markdown("---")
        
        # Section 2: Preprocessing Pipeline - Simple Explanation
        st.markdown("#### 📐 Image Preprocessing (Letterboxing)")
        st.markdown("""
        To preserve Sinhala character shapes, each image goes through:
        
        1. **Aspect Ratio Preservation** - Maintains original proportions
        2. **Resize to 384×384** - Standardized input size
        3. **White Padding** - Adds margins to prevent distortion
        4. **Normalization** - Prepares pixels for the model
        
        > **Why this matters:** Sinhala characters contain critical diacritics (පිලි) that get distorted with standard resizing.
        """)
        
        st.markdown("---")
        
        # Section 3: Core Model Configuration - Clean Metrics
        st.markdown("#### 📊 Model Configuration")
        
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.metric("Training Epochs", specs["training_epochs"])
        with col2:
            st.metric("Final Loss", f"{specs['training_loss']:.2f}")
        with col3:
            st.metric("Image Size", f"{specs['image_size']}×{specs['image_size']}px")
        with col4:
            st.metric("Max Length", f"{specs['max_length']} tokens")
        
        st.markdown("---")
        
        # Section 4: Generation Controls - Simple Cards
        st.markdown("#### 🔧 Generation Controls")
        
        control_col1, control_col2 = st.columns(2)
        with control_col1:
            st.markdown("""
            <div class="info-box">
                <strong>🎯 Repetition Penalty (2.0)</strong><br>
                Prevents character loops and repetitive patterns
            </div>
            """, unsafe_allow_html=True)
            
            st.markdown("""
            <div class="info-box">
                <strong>🔍 Beam Search (4-Beam)</strong><br>
                Explores multiple decoding paths for accuracy
            </div>
            """, unsafe_allow_html=True)
            
        with control_col2:
            st.markdown("""
            <div class="info-box">
                <strong>📏 Length Penalty (0.6)</strong><br>
                Discourages unnecessary extra characters
            </div>
            """, unsafe_allow_html=True)
            
            st.markdown("""
            <div class="info-box">
                <strong>⏹️ Early Stopping</strong><br>
                Stops generation when complete
            </div>
            """, unsafe_allow_html=True)
        
        st.markdown("""
        <div style="margin-top: 1rem; padding: 0.75rem; background: rgba(236, 185, 20, 0.05); border-radius: 8px;">
            <p style="margin: 0; color: #94A3B8; font-size: 13px;">
            <strong>💡 Additional Settings:</strong> `no_repeat_ngram_size=2` blocks duplicate patterns,
            ensuring clean, natural-looking Sinhala text output.
            </p>
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown('</div>', unsafe_allow_html=True)

    # FOOTER
    st.markdown("""
    <div style="text-align: center; padding: 20px; color: #475569; font-size: 12px; margin-top: 2rem; border-top: 1px solid #1E293B;">
        <p>© 2026 Sinhala Handwritten OCR | Powered by Fine-Tuned TrOCR | Intelligent Text Recognition System</p>
        <p style="margin-top: 5px;">Designed for Sinhala Handwritten Text Recognition | Version 2.0</p>
    </div>
    """, unsafe_allow_html=True)

# ==================== END OF APP ====================