import gradio as gr
import plotly.graph_objects as go
import plotly.express as px
from collections import Counter
import numpy as np

CODON_TABLE = {
    'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L',
    'CTT': 'L', 'CTC': 'L', 'CTA': 'L', 'CTG': 'L',
    'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M',
    'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
    'TCT': 'S', 'TCC': 'S', 'TCA': 'S', 'TCG': 'S',
    'CCT': 'P', 'CCC': 'P', 'CCA': 'P', 'CCG': 'P',
    'ACT': 'T', 'ACC': 'T', 'ACA': 'T', 'ACG': 'T',
    'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A',
    'TAT': 'Y', 'TAC': 'Y', 'TAA': '*', 'TAG': '*',
    'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
    'AAT': 'N', 'AAC': 'N', 'AAA': 'K', 'AAG': 'K',
    'GAT': 'D', 'GAC': 'D', 'GAA': 'E', 'GAG': 'E',
    'TGT': 'C', 'TGC': 'C', 'TGA': '*', 'TGG': 'W',
    'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
    'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
    'GGT': 'G', 'GGC': 'G', 'GGA': 'G', 'GGG': 'G',
}

COLOR_MAP = {'A': '#FF6B6B', 'T': '#4ECDC4', 'G': '#45B7D1', 'C': '#FFA07A'}

def analyze_dna(sequence):
    seq = sequence.upper().replace(" ", "").replace("\n", "")
    seq = ''.join(c for c in seq if c in 'ATGC')
    
    if len(seq) < 3:
        return "Enter at least 3 nucleotides", None, None, ""
    
    # GC Content sliding window
    window = min(50, len(seq) // 2) or 1
    gc_values = []
    for i in range(len(seq) - window + 1):
        w = seq[i:i+window]
        gc = (w.count('G') + w.count('C')) / len(w) * 100
        gc_values.append(gc)
    
    fig_gc = go.Figure()
    fig_gc.add_trace(go.Scatter(
        y=gc_values, mode='lines',
        fill='tozeroy', fillcolor='rgba(78,205,196,0.2)',
        line=dict(color='#4ECDC4', width=2)
    ))
    fig_gc.update_layout(
        title="GC Content (Sliding Window)",
        xaxis_title="Position", yaxis_title="GC %",
        template="plotly_dark", height=350,
        paper_bgcolor='#0f1117', plot_bgcolor='#0f1117'
    )
    
    # Nucleotide composition
    counts = Counter(seq)
    fig_comp = px.pie(
        names=list(counts.keys()),
        values=list(counts.values()),
        color=list(counts.keys()),
        color_discrete_map=COLOR_MAP,
        title="Nucleotide Composition"
    )
    fig_comp.update_layout(
        template="plotly_dark", height=350,
        paper_bgcolor='#0f1117', plot_bgcolor='#0f1117'
    )
    
    # Translate
    protein = []
    for i in range(0, len(seq) - 2, 3):
        codon = seq[i:i+3]
        aa = CODON_TABLE.get(codon, '?')
        protein.append(aa)
    protein_str = ''.join(protein)
    
    stats = f"""### 📊 Sequence Statistics
- **Length:** {len(seq)} bp
- **GC Content:** {(seq.count('G')+seq.count('C'))/len(seq)*100:.1f}%
- **A:** {counts.get('A',0)} | **T:** {counts.get('T',0)} | **G:** {counts.get('G',0)} | **C:** {counts.get('C',0)}
- **Protein:** `{protein_str[:60]}{'...' if len(protein_str)>60 else ''}`"""
    
    return stats, fig_gc, fig_comp, protein_str

demo = gr.Interface(
    fn=analyze_dna,
    inputs=gr.Textbox(label="DNA Sequence", lines=4,
        placeholder="Paste DNA sequence (e.g., ATGCGATCGATCG...)"),
    outputs=[
        gr.Markdown(label="Statistics"),
        gr.Plot(label="GC Content"),
        gr.Plot(label="Composition"),
        gr.Textbox(label="Protein Translation"),
    ],
    title="🧬 DNA Sequence Visualizer",
    description="Analyze and visualize DNA sequences with GC content, composition, and translation.",
    theme=gr.themes.Soft(),
)

if __name__ == "__main__":
    demo.launch()