import streamlit as st
import numpy as np
import pandas as pd
import re
from itertools import product
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from collections import Counter
import warnings
warnings.filterwarnings('ignore')

# === MULTI-SCALE KMER SETUP (3-mer, 4-mer, 5-mer) ===
ALL_KMERS_3 = [''.join(p) for p in product('ACGT', repeat=3)]
ALL_KMERS_4 = [''.join(p) for p in product('ACGT', repeat=4)]
ALL_KMERS_5 = [''.join(p) for p in product('ACGT', repeat=5)]

KMER_TO_INDEX_3 = {kmer: i for i, kmer in enumerate(ALL_KMERS_3)}
KMER_TO_INDEX_4 = {kmer: i for i, kmer in enumerate(ALL_KMERS_4)}
KMER_TO_INDEX_5 = {kmer: i for i, kmer in enumerate(ALL_KMERS_5)}

# === EXPANDED TF PATTERNS ===
TF_PATTERNS_EXPANDED = {
    'Dorsal-High': [re.compile(pattern) for pattern in ['GGGAAAACCC', 'GGGATTACCC', 'GGGAAATCCC']],
    'ERE-Half': [re.compile(pattern) for pattern in ['GGTCA', 'TGACC']],
    'p53-High': [re.compile(pattern) for pattern in ['[AG][AG][AG]CAT[CT][CT][CT]', '[AG][AG][AG]CTT[CT][CT][CT]']],
    'Dorsal-Low': [re.compile('[CT][CT][CT]A[AT][AT]CCC')],
    'TATA': [re.compile('TATAAA')],
    'CAAT': [re.compile('CCAAT')],
    'GC_Box': [re.compile('GGGCGG')],
    'AP1': [re.compile('TGA[GC]TCA')],
    'NF_kB': [re.compile('GGGACTTTCC')],
    'E2F': [re.compile('TTTCGCGC')],
    'SP1': [re.compile('GGGCGG')],
    'OCT': [re.compile('ATGCAAAT')]
}

# === LEGACY TF PATTERNS FOR FALLBACK ===
TF_PATTERNS = {
    'Dorsal-High': [re.compile(pattern) for pattern in ['GGGAAAACCC', 'GGGATTACCC', 'GGGAAATCCC']],
    'ERE-Half': [re.compile(pattern) for pattern in ['GGTCA', 'TGACC']],
    'p53-High': [re.compile(pattern) for pattern in ['[AG][AG][AG]CAT[CT][CT][CT]', '[AG][AG][AG]CTT[CT][CT][CT]']],
    'Dorsal-Low': [re.compile('[CT][CT][CT]A[AT][AT]CCC')]
}

# === LEGACY KMER SETUP FOR FALLBACK ===
ALL_KMERS = [''.join(p) for p in product('ACGT', repeat=3)]
KMER_TO_INDEX = {kmer: i for i, kmer in enumerate(ALL_KMERS)}

# === DEMO MODEL CLASS ===
class DemoModel:
    """Demo model for when actual models are not available"""
    def __init__(self):
        self.demo_mode = True
        self._params = 5650000
    
    def predict(self, X):
        """Generate demo predictions based on sequence features"""
        n_samples = X.shape[0]
        predictions = np.zeros((n_samples, 4))
        
        for i in range(n_samples):
            # Use some features to create realistic-looking predictions
            # Feature-based heuristics for demo
            feature_sum = np.sum(X[i])
            feature_mean = np.mean(X[i])
            feature_std = np.std(X[i]) if len(X[i]) > 1 else 1
            
            # Create probability distribution based on features
            if feature_mean > 10:  # High feature values
                predictions[i] = [0.1, 0.2, 0.6, 0.1]  # Favor super-enhancer
            elif feature_std < 1:  # Low variation
                predictions[i] = [0.7, 0.2, 0.05, 0.05]  # Favor negative
            elif feature_sum > 100:  # High total features
                predictions[i] = [0.1, 0.6, 0.2, 0.1]  # Favor normal
            else:
                predictions[i] = [0.2, 0.3, 0.3, 0.2]  # Balanced
            
            # Add some randomness
            noise = np.random.normal(0, 0.05, 4)
            predictions[i] += noise
            
            # Ensure probabilities sum to 1
            predictions[i] = np.abs(predictions[i])
            predictions[i] = predictions[i] / np.sum(predictions[i])
        
        return predictions
    
    def count_params(self):
        return self._params

class DemoScaler:
    """Demo scaler for when actual scaler is not available"""
    def __init__(self):
        self.demo_mode = True
    
    def transform(self, X):
        """Simple standardization for demo"""
        return (X - np.mean(X, axis=0)) / (np.std(X, axis=0) + 1e-8)

# === MODEL LOADING WITH HUGGING FACE SUPPORT ===
@st.cache_resource
def load_models():
    """Load models from Hugging Face with fallback to demo mode"""
    try:
        import joblib
        from tensorflow.keras.models import load_model
        import urllib.request
        import os
        
        # Hugging Face model URLs
        model_url = "https://huggingface.co/amitpande74/a100-beast-dna-enhancer/resolve/main/a100_beast_20250726_153354.h5"
        scaler_url = "https://huggingface.co/amitpande74/a100-beast-dna-enhancer/resolve/main/beast_scaler_20250726_153354.pkl"
        
        # Create models directory if it doesn't exist
        os.makedirs("models", exist_ok=True)
        
        model_path = "models/a100_beast_20250726_153354.h5"
        scaler_path = "models/beast_scaler_20250726_153354.pkl"
        
        # Download model if not exists
        if not os.path.exists(model_path):
            urllib.request.urlretrieve(model_url, model_path)
        
        # Download scaler if not exists
        if not os.path.exists(scaler_path):
            urllib.request.urlretrieve(scaler_url, scaler_path)
        
        # Load the models
        model = load_model(model_path)
        scaler = joblib.load(scaler_path)
        
        return model, scaler, "A100_Beast"
        
    except Exception as e:
        return DemoModel(), DemoScaler(), "Demo"

# Load models
model, scaler, model_type = load_models()

# === A100 BEAST FEATURE EXTRACTION ===
def extract_multi_kmer_features(sequences):
    """Extract multi-scale k-mer features"""
    X_3mer = np.zeros((len(sequences), len(ALL_KMERS_3)), dtype=np.float32)
    X_4mer = np.zeros((len(sequences), len(ALL_KMERS_4)), dtype=np.float32)
    X_5mer = np.zeros((len(sequences), len(ALL_KMERS_5)), dtype=np.float32)

    for seq_idx, sequence in enumerate(sequences):
        seq_len = len(sequence)
        
        # 3-mers
        if seq_len >= 3:
            for i in range(seq_len - 2):
                kmer = sequence[i:i+3]
                if kmer in KMER_TO_INDEX_3:
                    X_3mer[seq_idx, KMER_TO_INDEX_3[kmer]] += 1
        
        # 4-mers
        if seq_len >= 4:
            for i in range(seq_len - 3):
                kmer = sequence[i:i+4]
                if kmer in KMER_TO_INDEX_4:
                    X_4mer[seq_idx, KMER_TO_INDEX_4[kmer]] += 1
        
        # 5-mers
        if seq_len >= 5:
            for i in range(seq_len - 4):
                kmer = sequence[i:i+5]
                if kmer in KMER_TO_INDEX_5:
                    X_5mer[seq_idx, KMER_TO_INDEX_5[kmer]] += 1

    return X_3mer, X_4mer, X_5mer

def calculate_comprehensive_palindrome_features(sequences):
    """Calculate palindrome features"""
    features = np.zeros((len(sequences), 35), dtype=np.float32)
    complement = str.maketrans('ATGC', 'TACG')

    for seq_idx, sequence in enumerate(sequences):
        seq = sequence.replace('N', '')
        seq_len = len(seq)

        if seq_len < 4:
            continue

        # Palindrome counters
        pal_4 = pal_5 = pal_6 = pal_7 = pal_8 = pal_9 = 0
        pal_10_15 = pal_16_plus = 0
        imp_1mm = imp_2mm = imp_3mm = 0
        
        # TF counters
        tf_counts = {name: 0 for name in TF_PATTERNS_EXPANDED.keys()}

        # Palindrome detection (4-25 bp range)
        max_check = min(seq_len, 25)
        for length in range(4, max_check + 1):
            for start in range(seq_len - length + 1):
                substr = seq[start:start + length]
                rev_comp = substr.translate(complement)[::-1]

                if substr == rev_comp:  # Perfect palindrome
                    if length == 4: pal_4 += 1
                    elif length == 5: pal_5 += 1
                    elif length == 6: pal_6 += 1
                    elif length == 7: pal_7 += 1
                    elif length == 8: pal_8 += 1
                    elif length == 9: pal_9 += 1
                    elif 10 <= length <= 15: pal_10_15 += 1
                    else: pal_16_plus += 1
                else:  # Imperfect palindromes
                    mismatches = sum(a != b for a, b in zip(substr, rev_comp))
                    if mismatches == 1: imp_1mm += 1
                    elif mismatches == 2: imp_2mm += 1
                    elif mismatches == 3: imp_3mm += 1

        # TF binding site detection
        for tf_name, patterns in TF_PATTERNS_EXPANDED.items():
            for pattern in patterns:
                tf_counts[tf_name] += len(pattern.findall(seq))

        # Calculate metrics
        total_perfect = pal_4 + pal_5 + pal_6 + pal_7 + pal_8 + pal_9 + pal_10_15 + pal_16_plus
        total_imperfect = imp_1mm + imp_2mm + imp_3mm
        total_tf = sum(tf_counts.values())
        
        # Shadow detection metrics
        perfect_imperfect_ratio = total_perfect / max(total_imperfect, 1)
        imperfect_perfect_ratio = total_imperfect / max(total_perfect, 1)
        
        shadow_flag_strict = 1.0 if perfect_imperfect_ratio < 0.1 else 0.0
        shadow_flag_loose = 1.0 if perfect_imperfect_ratio < 0.2 else 0.0
        imperfect_dominance = 1.0 if total_imperfect > total_perfect * 6 else 0.0
        
        # Densities
        imperfect_density = total_imperfect / max(seq_len / 1000, 0.1)
        perfect_density = total_perfect / max(seq_len / 1000, 0.1)
        tf_density = total_tf / max(seq_len / 1000, 0.1)
        
        # Ratios
        short_long_ratio = (pal_4 + pal_5 + pal_6) / max(pal_10_15 + pal_16_plus, 1)
        mismatch_gradient = imp_1mm / max(imp_1mm + imp_2mm + imp_3mm, 1)
        
        # Feature vector (35 features)
        features[seq_idx] = [
            pal_4, pal_5, pal_6, pal_7, pal_8, pal_9, pal_10_15, pal_16_plus,
            imp_1mm, imp_2mm, imp_3mm,
            total_perfect, total_imperfect, total_tf,
            perfect_imperfect_ratio, imperfect_perfect_ratio, short_long_ratio, mismatch_gradient, tf_density,
            shadow_flag_strict, shadow_flag_loose, imperfect_dominance, imperfect_density,
            tf_counts['Dorsal-High'], tf_counts['ERE-Half'], tf_counts['p53-High'], 
            tf_counts['Dorsal-Low'], tf_counts['TATA'], tf_counts['CAAT'],
            tf_counts['GC_Box'], tf_counts['AP1'], tf_counts['NF_kB'],
            tf_counts['E2F'], tf_counts['SP1'], tf_counts['OCT']
        ]

    return features

def calculate_advanced_basic_features(sequences):
    """Calculate basic sequence features"""
    features = np.zeros((len(sequences), 12), dtype=np.float32)

    for i, seq in enumerate(sequences):
        length = len(seq)
        if length > 0:
            a_freq = seq.count('A') / length
            t_freq = seq.count('T') / length
            g_freq = seq.count('G') / length
            c_freq = seq.count('C') / length
            
            gc_content = g_freq + c_freq
            at_content = a_freq + t_freq
            
            gc_skew = (g_freq - c_freq) / max(gc_content, 0.001)
            at_skew = (a_freq - t_freq) / max(at_content, 0.001)
            
            cpg_density = seq.count('CG') / max(length - 1, 1)
            cpa_density = seq.count('CA') / max(length - 1, 1)
            tpg_density = seq.count('TG') / max(length - 1, 1)
            
            purine_content = a_freq + g_freq
            
        else:
            a_freq = t_freq = g_freq = c_freq = gc_content = at_content = 0
            gc_skew = at_skew = cpg_density = cpa_density = tpg_density = purine_content = 0

        features[i] = [
            length, gc_content, at_content, gc_skew, at_skew,
            cpg_density, cpa_density, tpg_density, purine_content,
            a_freq, t_freq, g_freq
        ]

    return features

def extract_features_a100_beast(sequences):
    """Extract ALL features exactly like A100 Beast training (1391 features)"""
    # Multi-scale k-mer features
    X_3mer, X_4mer, X_5mer = extract_multi_kmer_features(sequences)
    
    # Comprehensive palindrome features
    palindrome_features = calculate_comprehensive_palindrome_features(sequences)
    
    # Advanced basic features
    basic_features = calculate_advanced_basic_features(sequences)

    # Apply Shadow boosting
    shadow_boost = 5.0
    palindrome_features[:, [14, 15, 19, 20, 21, 22]] *= shadow_boost
    
    # Combine features: 64 + 256 + 1024 + 35 + 12 = 1391 features
    X = np.hstack([X_3mer, X_4mer, X_5mer, palindrome_features, basic_features])
    
    return X.astype(np.float32)

# === LEGACY FEATURE EXTRACTION (FALLBACK) ===
def extract_features(sequences):
    """Extract features for legacy ML model."""
    X_kmers = np.zeros((len(sequences), len(ALL_KMERS)), dtype=np.float32)
    X_basic = np.zeros((len(sequences), 4), dtype=np.float32)
    X_pal = np.zeros((len(sequences), 17), dtype=np.float32)
    
    complement = str.maketrans('ATGC', 'TACG')

    for i, seq in enumerate(sequences):
        seq = seq.upper().replace("N", "")
        L = len(seq)

        # k-mer features
        for j in range(len(seq) - 2):
            kmer = seq[j:j+3]
            if kmer in KMER_TO_INDEX:
                X_kmers[i, KMER_TO_INDEX[kmer]] += 1

        # basic features
        gc_content = (seq.count('G') + seq.count('C')) / L if L else 0
        at_content = (seq.count('A') + seq.count('T')) / L if L else 0
        cpg_density = seq.count('CG') / max(L - 1, 1)
        X_basic[i] = [L, gc_content, at_content, cpg_density]

        # palindrome features
        short_pal = medium_pal = long_pal = 0
        mild_imp = moderate_imp = 0
        high_tf = low_tf = 0

        for length in range(4, min(L, 20) + 1):
            for start in range(L - length + 1):
                substr = seq[start:start + length]
                rev_comp = substr.translate(complement)[::-1]
                if substr == rev_comp:
                    if length <= 5:
                        short_pal += 1
                    elif length <= 9:
                        medium_pal += 1
                    else:
                        long_pal += 1
                else:
                    mismatches = sum(a != b for a, b in zip(substr, rev_comp))
                    if mismatches == 1:
                        mild_imp += 1
                    elif mismatches == 2:
                        moderate_imp += 1

        for tf_type, patterns in TF_PATTERNS.items():
            for pattern in patterns:
                matches = len(pattern.findall(seq))
                if 'High' in tf_type:
                    high_tf += matches
                else:
                    low_tf += matches

        total_perfect = short_pal + medium_pal + long_pal
        total_imperfect = mild_imp + moderate_imp
        total_tf = high_tf + low_tf

        X_pal[i] = [
            short_pal, medium_pal, long_pal, total_perfect,
            mild_imp, moderate_imp, total_imperfect,
            total_perfect / max(total_imperfect, 1),
            long_pal / max(short_pal, 1),
            total_perfect / max(L / 100, 1),
            high_tf, low_tf, total_tf,
            total_tf / max(L / 100, 1),
            long_pal * 6,
            (long_pal * 20 + medium_pal * 10 + high_tf * 15 + (total_perfect / max(total_imperfect, 1)) * 10),
            (total_imperfect * 12 + low_tf * 8 + (1 / max(total_perfect / max(total_imperfect, 1), 0.1)) * 10)
        ]

    return np.hstack([X_kmers, X_pal, X_basic])

# === HELPER FUNCTIONS ===
def calculate_palindrome_ratio(sequence):
    """Calculate perfect/imperfect palindrome ratio"""
    complement = str.maketrans('ATGC', 'TACG')
    total_perfect = total_imperfect = 0
    
    for length in range(4, min(len(sequence), 20) + 1):
        for start in range(len(sequence) - length + 1):
            substr = sequence[start:start + length]
            rev_comp = substr.translate(complement)[::-1]
            if substr == rev_comp:
                total_perfect += 1
            else:
                mismatches = sum(a != b for a, b in zip(substr, rev_comp))
                if mismatches <= 2:
                    total_imperfect += 1
    
    return total_perfect / max(total_imperfect, 1)

def parse_fasta(content):
    """Parse FASTA content and return list of sequences with their headers."""
    sequences = []
    headers = []
    current_header = ""
    current_sequence = ""
    
    lines = content.strip().splitlines()
    
    for line in lines:
        line = line.strip()
        if line.startswith(">"):
            # Save previous sequence if it exists
            if current_sequence:
                sequences.append(current_sequence)
                headers.append(current_header)
            # Start new sequence
            current_header = line[1:] if line[1:] else f"Sequence_{len(sequences)+1}"
            current_sequence = ""
        elif line:  # Only process non-empty lines
            # Remove any non-DNA characters and concatenate
            clean_line = re.sub(r'[^ATCGN]', '', line.upper())
            current_sequence += clean_line
    
    # Don't forget the last sequence
    if current_sequence:
        sequences.append(current_sequence)
        headers.append(current_header)
    
    return sequences, headers

def analyze_composition(sequence):
    """Analyze nucleotide composition and patterns."""
    length = len(sequence)
    if length == 0:
        return {
            'composition': {'A': 0, 'T': 0, 'G': 0, 'C': 0},
            'gc_content': 0, 'at_content': 0, 'gc_skew': 0, 'at_skew': 0,
            'purine_content': 0, 'pyrimidine_content': 0, 'length': 0
        }
    
    composition = {
        'A': sequence.count('A') / length,
        'T': sequence.count('T') / length,
        'G': sequence.count('G') / length,
        'C': sequence.count('C') / length
    }
    
    # Calculate various metrics
    gc_content = composition['G'] + composition['C']
    at_content = composition['A'] + composition['T']
    gc_skew = (composition['G'] - composition['C']) / (composition['G'] + composition['C']) if (composition['G'] + composition['C']) > 0 else 0
    at_skew = (composition['A'] - composition['T']) / (composition['A'] + composition['T']) if (composition['A'] + composition['T']) > 0 else 0
    
    # Purine/Pyrimidine content
    purine_content = composition['A'] + composition['G']  # A, G
    pyrimidine_content = composition['T'] + composition['C']  # T, C
    
    return {
        'composition': composition,
        'gc_content': gc_content,
        'at_content': at_content,
        'gc_skew': gc_skew,
        'at_skew': at_skew,
        'purine_content': purine_content,
        'pyrimidine_content': pyrimidine_content,
        'length': length
    }

def find_palindromes(sequence, min_length=4, max_length=20):
    """Find all palindromes in a sequence."""
    palindromes = []
    complement = str.maketrans('ATGC', 'TACG')
    
    for length in range(min_length, min(len(sequence), max_length) + 1):
        for start in range(len(sequence) - length + 1):
            substr = sequence[start:start + length]
            rev_comp = substr.translate(complement)[::-1]
            if substr == rev_comp:
                palindromes.append({
                    'sequence': substr,
                    'start': start,
                    'end': start + length,
                    'length': length,
                    'type': 'perfect'
                })
    
    return palindromes

def find_imperfect_palindromes(sequence, min_length=4, max_length=20, max_mismatches=2):
    """Find imperfect palindromes with up to max_mismatches."""
    imperfect_palindromes = []
    complement = str.maketrans('ATGC', 'TACG')
    
    for length in range(min_length, min(len(sequence), max_length) + 1):
        for start in range(len(sequence) - length + 1):
            substr = sequence[start:start + length]
            rev_comp = substr.translate(complement)[::-1]
            
            mismatches = sum(a != b for a, b in zip(substr, rev_comp))
            if 0 < mismatches <= max_mismatches:
                imperfect_palindromes.append({
                    'sequence': substr,
                    'start': start,
                    'end': start + length,
                    'length': length,
                    'mismatches': mismatches,
                    'type': f'{mismatches}_mismatch'
                })
    
    return imperfect_palindromes

# === STREAMLIT APP ===
st.set_page_config(page_title="A100 Beast DNA Enhancer Analyzer", layout="wide")

# Header
st.title("🔥 A100 Beast DNA Enhancer Analyzer")
st.markdown("Advanced DNA sequence analysis with A100 Beast ML model featuring 1391 features for precise enhancer classification.")

# Display model info in sidebar
st.sidebar.header("🤖 Model Information")
if model is not None:
    if model_type == "Demo":
        st.sidebar.warning("🔄 Demo Mode")
        st.sidebar.info("Running in demo mode with simulated predictions")
    else:
        st.sidebar.success(f"✅ Model: {model_type}")
    
    params = model.count_params() if hasattr(model, 'count_params') else "Unknown"
    st.sidebar.info(f"Parameters: {params:,}")
    
    if model_type == "A100_Beast":
        st.sidebar.success("🔥 A100 Beast Features: 1391")
        st.sidebar.write("- 64 3-mer k-mers")
        st.sidebar.write("- 256 4-mer k-mers") 
        st.sidebar.write("- 1024 5-mer k-mers")
        st.sidebar.write("- 35 palindrome features")
        st.sidebar.write("- 12 basic features")
    elif model_type == "Demo":
        st.sidebar.info("Demo Features: Simulated")
        st.sidebar.write("- Feature-based heuristics")
        st.sidebar.write("- Randomized predictions")
        st.sidebar.write("- Educational purposes")
    else:
        st.sidebar.info(f"Features: {85 if model_type == 'Enhanced' else 'Variable'}")
else:
    st.sidebar.error("🚫 No model available")

# Sidebar for settings
st.sidebar.header("⚙️ Analysis Settings")

analysis_options = st.sidebar.multiselect(
    "Select Analysis Types:",
    ["Enhancer Classification", "Sequence Composition", "Palindrome Analysis"],
    default=["Enhancer Classification", "Sequence Composition", "Palindrome Analysis"]
)

if model_type in ["A100_Beast", "Demo"]:
    st.sidebar.subheader("🔥 A100 Beast Options")
    show_feature_analysis = st.sidebar.checkbox("Show detailed feature analysis", value=True)
    show_palindrome_details = st.sidebar.checkbox("Show palindrome breakdown", value=True)

palindrome_min_length = st.sidebar.slider("Minimum Palindrome Length", 4, 10, 4)
palindrome_max_length = st.sidebar.slider("Maximum Palindrome Length", 10, 30, 20)

# Input section
with st.expander("ℹ️ How to use A100 Beast Analyzer", expanded=False):
    st.markdown("""
    **A100 Beast Model Features:**
    - 🔥 **1391 Advanced Features**: Multi-scale k-mer analysis (3-mer, 4-mer, 5-mer)
    - 🎯 **4-Class Classification**: Negative, Normal, Super-enhancer, Shadow
    - 🧬 **Comprehensive Analysis**: 35 palindrome features + 12 basic features
    - 📊 **Shadow Boosting**: Enhanced detection of shadow enhancers
    - 🚀 **5.65M Parameters**: Deep neural network for complex pattern recognition
    
    **Input Formats:**
    ```
    >Shadow_enhancer
    CAAGCAAGTTTCAGCTCCCACTGCCGCCCCTCCGGCC...
    >Normal_enhancer  
    TATGGTTGTGCTTTTTTTTTTTCTTTAAGAGAAAT...
    ```
    """)

col1, col2 = st.columns([2, 1])

with col1:
    uploaded_file = st.file_uploader("📂 Upload FASTA file", type=["fasta", "fa", "txt"])
    seq_input = st.text_area("📋 Or paste sequences:", height=150, 
                            placeholder="Paste your DNA sequences here...")

with col2:
    st.markdown("**Quick Examples:**")
    if st.button("🧪 Load A100 Beast Test Sequences"):
        example_seqs = """>Shadow_enhancer
CAAGCAAGTTTCAGCTCCCACTGCCGCCCCTCCGGCCAAGGTCATCGGCCAGCCTGTGCTCGCTGCTCCAGGCAAAAATTCTTCCAATTCCAGCAGCACAACTGAGTGGTAAGTGGCACTGCACTGGCAGAAATTGCCGCATCTGGAAACTTAAAGTCTAGTCCATTTGAGCAGTTTTAAAGCTGCAATCCCTTTCGTCCCACTGGATTCTACTACTTTAAATTCAATTTACGTTGTTTGCGTAGGCTTCAAAGAATCATAGTTCATGTATTGCACAA
>Normal_enhancer
TATGGTTGTGCTTTTTTTTTTTCTTTAAGAGAAATAAAGAAAGCACATAAAAAGACCCTCTCCTTCACTTACAGGGTTTCTCACTCCCCGCCAGGGCACATAGCCCCAGGACGAGGAGCGCTGCCAGGAGCTGGGCGTCTCGCCGTCCCATGTCTAGCTCAGCTGCACCCCAGGGTGGCTTGCAGAATGCATGGTGTCCACTGCCGGGTATGTTTTATAATCTCCCCTCTGTTTGTCCAA
>Negative_sequence
CTACCTGGGACCTTCGATGCGTTCCAATAAATTGTGCCTTATGTAGTTATCGGATAGCTATGTGGCGTTGCTAAATTTAGGAGTCTTGGTCACAAGTTGAGCCGATGGTCAGGCTCCGTTAAGGATGCGCGTACGCGGGACCGATAGTAGGGAACGTATGCTTTGATAATCGCTAGTTGACAGTGTGCTGGTCTATTGAGTTTCCTCGCTAAAATTTCGGGCCTTCGAGTTTC"""
        st.session_state.seq_input = example_seqs

# Process sequences when button is clicked
if st.button("🚀 Analyze Sequences", type="primary"):
    sequences = []
    headers = []
    
    # Parse input
    if uploaded_file:
        content = uploaded_file.read().decode("utf-8")
        if content.strip().startswith(">"):
            sequences, headers = parse_fasta(content)
            st.success(f"📄 Parsed FASTA file: Found {len(sequences)} sequences")
        else:
            content_lines = [line.strip().upper() for line in content.strip().splitlines() 
                           if line.strip() and not line.strip().startswith(">")]
            concatenated = ''.join([re.sub(r'[^ATGC]', '', line) for line in content_lines])
            sequences = [concatenated]
            headers = ["Sequence_1"]
            st.success(f"📄 Parsed as single sequence: {len(sequences[0])} bp")
    
    elif seq_input.strip():
        input_text = seq_input.strip()
        
        if input_text.startswith(">"):
            sequences, headers = parse_fasta(input_text)
            st.info(f"📝 Parsed as FASTA format: Found {len(sequences)} sequences")
        else:
            input_lines = [line.strip().upper() for line in input_text.splitlines() if line.strip()]
            concatenated = ''.join([re.sub(r'[^ATGC]', '', line) for line in input_lines])
            sequences = [concatenated]
            headers = ["Sequence_1"]
            st.info(f"📝 Single sequence: {len(sequences[0])} bp")
    
    # Validate sequences
    valid_sequences = []
    valid_headers = []
    
    for i, seq in enumerate(sequences):
        clean_seq = re.sub(r'[^ATCGN]', '', seq.upper())
        if len(clean_seq) >= 10:
            valid_sequences.append(clean_seq)
            valid_headers.append(headers[i] if i < len(headers) else f"Sequence_{i+1}")
        else:
            st.warning(f"⚠️ Skipped sequence {i+1}: Too short or invalid")

    if len(valid_sequences) == 0:
        st.error("❗ No valid DNA sequences found.")
    else:
        # Create tabs for different analyses
        tabs = st.tabs(["📊 Overview"] + [f"🧬 {header[:20]}..." if len(header) > 20 else f"🧬 {header}" for header in valid_headers[:3]])
        
        with tabs[0]:  # Overview tab
            col1, col2, col3, col4 = st.columns(4)
            
            with col1:
                st.metric("Total Sequences", len(valid_sequences))
            with col2:
                avg_length = np.mean([len(seq) for seq in valid_sequences])
                st.metric("Average Length", f"{avg_length:.0f} bp")
            with col3:
                total_length = sum(len(seq) for seq in valid_sequences)
                st.metric("Total Length", f"{total_length:,} bp")
            with col4:
                avg_gc = np.mean([analyze_composition(seq)['gc_content'] for seq in valid_sequences])
                st.metric("Average GC%", f"{avg_gc:.1%}")
            
            # Enhanced Classification
            if "Enhancer Classification" in analysis_options and model and scaler:
                with st.spinner('🔄 Running classification...'):
                    try:
                        # Use appropriate feature extraction
                        if model_type in ["A100_Beast", "Demo"]:
                            X = extract_features_a100_beast(valid_sequences)
                            if model_type == "A100_Beast":
                                st.success(f"🔥 A100 Beast: Extracted {X.shape[1]} features")
                            else:
                                st.info(f"🔄 Demo Mode: Extracted {X.shape[1]} features")
                        else:
                            X = extract_features(valid_sequences)
                            st.info(f"Enhanced model: Extracted {X.shape[1]} features")
                        
                        # Scale and predict
                        X_scaled = scaler.transform(X)
                        preds = model.predict(X_scaled)
                        pred_labels = np.argmax(preds, axis=1)
                        pred_probs = np.max(preds, axis=1)
                        
                        class_map = {0: 'Negative', 1: 'Normal', 2: 'Super-enhancer', 3: 'Shadow'}
                        pred_names = [class_map[i] for i in pred_labels]
                        
                        # Results
                        st.subheader("🎯 Classification Results")
                        
                        if model_type == "Demo":
                            st.warning("⚠️ Demo predictions shown - install TensorFlow and joblib for real model")
                        
                        # Model metrics
                        col1, col2, col3, col4 = st.columns(4)
                        with col1:
                            st.metric("Model", model_type)
                        with col2:
                            st.metric("Features", f"{X.shape[1]:,}")
                        with col3:
                            params = model.count_params() if hasattr(model, 'count_params') else "Unknown"
                            st.metric("Parameters", f"{params:,}")
                        with col4:
                            avg_conf = np.mean(pred_probs)
                            st.metric("Avg Confidence", f"{avg_conf:.1%}")
                        
                        # Results table
                        results_df = pd.DataFrame({
                            'Sequence': valid_headers,
                            'Length': [len(seq) for seq in valid_sequences],
                            'Classification': pred_names,
                            'Confidence': [f"{prob:.1%}" for prob in pred_probs],
                            'Palindrome_Ratio': [f"{calculate_palindrome_ratio(seq):.3f}" for seq in valid_sequences]
                        })
                        
                        st.dataframe(results_df, use_container_width=True)
                        
                        # Detailed probabilities
                        st.subheader("🔍 Detailed Probabilities")
                        prob_df = pd.DataFrame(preds, columns=['Negative', 'Normal', 'Super-enhancer', 'Shadow'])
                        prob_df['Sequence'] = valid_headers
                        prob_df = prob_df[['Sequence', 'Negative', 'Normal', 'Super-enhancer', 'Shadow']]
                        
                        for col in ['Negative', 'Normal', 'Super-enhancer', 'Shadow']:
                            prob_df[col] = prob_df[col].apply(lambda x: f"{x:.1%}")
                        
                        st.dataframe(prob_df, use_container_width=True)
                        
                        # Visualization
                        col1, col2 = st.columns(2)
                        
                        with col1:
                            class_counts = pd.Series(pred_names).value_counts()
                            fig_pie = px.pie(values=class_counts.values, names=class_counts.index,
                                           title="Classification Results", 
                                           color_discrete_sequence=px.colors.qualitative.Set3)
                            st.plotly_chart(fig_pie, use_container_width=True)
                        
                        with col2:
                            fig_conf = px.histogram(x=pred_probs, nbins=10, title="Confidence Distribution",
                                                  labels={'x': 'Confidence Score', 'y': 'Count'})
                            st.plotly_chart(fig_conf, use_container_width=True)
                        
                    except Exception as e:
                        st.error(f"❌ Classification error: {str(e)}")
                        import traceback
                        st.code(traceback.format_exc())
            
            # Sequence composition overview
            if "Sequence Composition" in analysis_options:
                st.subheader("📈 Composition Summary")
                all_compositions = [analyze_composition(seq) for seq in valid_sequences]
                
                comp_df = pd.DataFrame({
                    'Sequence': valid_headers,
                    'Length': [comp['length'] for comp in all_compositions],
                    'GC%': [f"{comp['gc_content']:.1%}" for comp in all_compositions],
                    'AT%': [f"{comp['at_content']:.1%}" for comp in all_compositions],
                    'GC Skew': [f"{comp['gc_skew']:.3f}" for comp in all_compositions],
                    'AT Skew': [f"{comp['at_skew']:.3f}" for comp in all_compositions]
                })
                
                st.dataframe(comp_df, use_container_width=True)
        
        # Individual sequence tabs (limit to first 3)
        for i, (seq, header) in enumerate(zip(valid_sequences[:3], valid_headers[:3])):
            with tabs[i+1]:
                st.subheader(f"Analysis: {header}")
                st.code(f"Length: {len(seq)} bp")
                
                # Show ML prediction if classification was done
                if "Enhancer Classification" in analysis_options and model and scaler and 'pred_names' in locals():
                    prediction = pred_names[i]
                    confidence = pred_probs[i]
                    st.success(f"🎯 **A100 Beast Prediction: {prediction}** (Confidence: {confidence:.1%})")
                
                if len(seq) > 200:
                    st.text_area("Sequence Preview (first 200 bp):", seq[:200] + "...", height=60, key=f"seq_preview_{i}")
                else:
                    st.text_area("Full Sequence:", seq, height=60, key=f"seq_full_{i}")
                
                # Palindrome Analysis
                if "Palindrome Analysis" in analysis_options:
                    with st.expander("🔄 Palindrome Analysis", expanded=True):
                        palindromes = find_palindromes(seq, palindrome_min_length, palindrome_max_length)
                        imperfect_palindromes = find_imperfect_palindromes(seq, palindrome_min_length, palindrome_max_length)
                        
                        col1, col2, col3, col4 = st.columns(4)
                        with col1:
                            st.metric("Perfect Palindromes", len(palindromes))
                        with col2:
                            st.metric("Imperfect Palindromes", len(imperfect_palindromes))
                        with col3:
                            perfect_6plus = len([p for p in palindromes if p['length'] >= 6])
                            st.metric("Perfect (≥6bp)", perfect_6plus)
                        with col4:
                            ratio = len(palindromes) / max(len(imperfect_palindromes), 1)
                            st.metric("Perfect/Imperfect Ratio", f"{ratio:.2f}")
                        
                        # Show ML prediction context, not just ratio interpretation
                        if "Enhancer Classification" in analysis_options and model and scaler and 'pred_names' in locals():
                            prediction = pred_names[i]
                            if prediction == "Shadow":
                                st.warning(f"🟡 **A100 Beast detected Shadow enhancer** - Ratio: {ratio:.3f}")
                                st.info("The ML model considers 1391 features beyond just palindrome ratios for classification.")
                            elif prediction == "Super-enhancer":
                                st.success(f"🟢 **A100 Beast detected Super-enhancer** - Strong regulatory potential")
                            elif prediction == "Normal":
                                st.info(f"🔵 **A100 Beast detected Normal enhancer** - Standard regulatory element")
                            else:
                                st.error(f"🔴 **A100 Beast detected Non-enhancer** - No regulatory activity predicted")
                        else:
                            # Fallback to simple ratio interpretation if no ML prediction
                            if ratio < 0.15:
                                st.warning(f"🟡 **Shadow-like ratio detected**: {ratio:.3f} < 0.15 suggests shadow enhancer characteristics")
                            elif ratio > 0.5:
                                st.success(f"🟢 **High palindrome ratio**: {ratio:.3f} suggests normal enhancer patterns")
                            else:
                                st.info(f"🔵 **Moderate ratio**: {ratio:.3f} - mixed palindrome patterns")

# Footer
st.markdown("---")
st.markdown(f"""
**🔥 A100 Beast DNA Enhancer Analyzer** - Advanced ML-powered sequence analysis  
**Model Status:** {model_type} {"(Demo Mode)" if model_type == "Demo" else ""}
- **Deep Neural Network**: Advanced architecture for enhancer classification
- **Multi-scale Features**: K-mer analysis + palindrome patterns + sequence composition  
- **4-Class Prediction**: Negative, Normal, Super-enhancer, Shadow with confidence scores
- **Shadow Detection**: Enhanced identification of shadow enhancers
- **Real-time Analysis**: Instant classification and detailed feature breakdown

{f"**Note:** Running in demo mode. For full functionality, ensure TensorFlow and model files are available." if model_type == "Demo" else ""}
""")