import gradio as gr
import torch
import joblib
import numpy as np
from itertools import product
import torch.nn as nn
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import seaborn as sns
from PIL import Image
import io
import pandas as pd
from typing import Tuple, List, Dict, Any
from dataclasses import dataclass
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots

###############################################################################
# 1. DATA STRUCTURES & MODEL
###############################################################################

@dataclass
class SequenceAnalysis:
    """Container for sequence analysis results"""
    header: str
    sequence: str
    length: int
    gc_content: float
    classification: str
    human_prob: float
    nonhuman_prob: float
    shap_values: np.ndarray
    shap_means: np.ndarray
    extreme_regions: Dict[str, Dict[str, Any]]

class VirusClassifier(nn.Module):
    def __init__(self, input_shape: int):
        super(VirusClassifier, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_shape, 64),
            nn.GELU(),
            nn.BatchNorm1d(64),
            nn.Dropout(0.3),
            nn.Linear(64, 32),
            nn.GELU(),
            nn.BatchNorm1d(32),
            nn.Dropout(0.3),
            nn.Linear(32, 32),
            nn.GELU(),
            nn.Linear(32, 2)
        )
        
    def forward(self, x):
        return self.network(x)

###############################################################################
# 2. SEQUENCE PROCESSING
###############################################################################

def parse_fasta(text: str) -> List[Tuple[str, str]]:
    """Parse FASTA formatted text with improved robustness"""
    sequences = []
    current_header = None
    current_sequence = []
    
    for line in text.strip().split('\n'):
        line = line.strip()
        if not line:
            continue
        if line.startswith('>'):
            if current_header:
                sequences.append((current_header, ''.join(current_sequence)))
            current_header = line[1:]
            current_sequence = []
        else:
            # Filter out non-ACGT characters and convert to uppercase
            filtered_line = ''.join(c for c in line.upper() if c in 'ACGT')
            current_sequence.append(filtered_line)
            
    if current_header:
        sequences.append((current_header, ''.join(current_sequence)))
    return sequences

def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
    """Convert sequence to k-mer frequency vector with optimizations"""
    kmers = [''.join(p) for p in product("ACGT", repeat=k)]
    kmer_dict = {km: i for i, km in enumerate(kmers)}
    vec = np.zeros(len(kmers), dtype=np.float32)
    
    # Use sliding window for efficiency
    for i in range(len(sequence) - k + 1):
        kmer = sequence[i:i+k]
        if kmer in kmer_dict:  # Handle non-ACGT kmers
            vec[kmer_dict[kmer]] += 1
    
    # Normalize
    total_kmers = len(sequence) - k + 1
    if total_kmers > 0:
        vec = vec / total_kmers
    
    return vec

def compute_gc_content(sequence: str) -> float:
    """Compute GC content percentage"""
    if not sequence:
        return 0.0
    gc_count = sum(1 for base in sequence if base in 'GC')
    return (gc_count / len(sequence)) * 100.0

###############################################################################
# 3. SHAP & ANALYSIS
###############################################################################

def calculate_shap_values(model: nn.Module, x_tensor: torch.Tensor) -> Tuple[np.ndarray, float]:
    """Calculate SHAP values using ablation with improved efficiency"""
    model.eval()
    with torch.no_grad():
        baseline_output = model(x_tensor)
        baseline_probs = torch.softmax(baseline_output, dim=1)
        baseline_prob = baseline_probs[0, 1].item()
        
        shap_values = []
        x_zeroed = x_tensor.clone()
        
        # Vectorized computation where possible
        for i in range(x_tensor.shape[1]):
            x_zeroed[0, i] = 0.0
            output = model(x_zeroed)
            probs = torch.softmax(output, dim=1)
            impact = baseline_prob - probs[0, 1].item()
            shap_values.append(impact)
            x_zeroed[0, i] = x_tensor[0, i]
            
    return np.array(shap_values), baseline_prob

def compute_positionwise_scores(sequence: str, shap_values: np.ndarray, k: int = 4) -> np.ndarray:
    """Compute per-base SHAP scores with optimized memory usage"""
    kmers = [''.join(p) for p in product("ACGT", repeat=k)]
    kmer_dict = {km: i for i, km in enumerate(kmers)}
    
    seq_len = len(sequence)
    shap_sums = np.zeros(seq_len, dtype=np.float32)
    coverage = np.zeros(seq_len, dtype=np.float32)
    
    # Vectorized operations where possible
    for i in range(seq_len - k + 1):
        kmer = sequence[i:i+k]
        if kmer in kmer_dict:
            idx = kmer_dict[kmer]
            shap_sums[i:i+k] += shap_values[idx]
            coverage[i:i+k] += 1
    
    with np.errstate(divide='ignore', invalid='ignore'):
        shap_means = np.where(coverage > 0, shap_sums / coverage, 0.0)
    
    return shap_means

def find_extreme_regions(shap_means: np.ndarray, window_size: int = 500) -> Dict[str, Dict[str, Any]]:
    """Find regions with extreme SHAP values using efficient sliding window"""
    if len(shap_means) < window_size:
        window_size = len(shap_means)
    
    # Compute cumulative sum for efficient sliding window
    cumsum = np.cumsum(np.pad(shap_means, (0, 1)))
    
    # Sliding window calculation
    window_avgs = (cumsum[window_size:] - cumsum[:-window_size]) / window_size
    
    max_idx = np.argmax(window_avgs)
    min_idx = np.argmin(window_avgs)
    
    return {
        "human": {
            "start": max_idx,
            "end": max_idx + window_size,
            "avg_shap": float(window_avgs[max_idx])
        },
        "nonhuman": {
            "start": min_idx,
            "end": min_idx + window_size,
            "avg_shap": float(window_avgs[min_idx])
        }
    }

###############################################################################
# 4. VISUALIZATION
###############################################################################

def create_genome_overview_plot(analysis: SequenceAnalysis) -> go.Figure:
    """Create an interactive genome overview using Plotly"""
    fig = make_subplots(
        rows=2, cols=1,
        subplot_titles=("SHAP Values Along Genome", "GC Content"),
        row_heights=[0.7, 0.3],
        vertical_spacing=0.1
    )
    
    # SHAP trace
    fig.add_trace(
        go.Scatter(
            x=list(range(len(analysis.shap_means))),
            y=analysis.shap_means,
            name="SHAP",
            line=dict(color='rgba(31, 119, 180, 0.8)'),
            hovertemplate="Position: %{x}<br>SHAP: %{y:.4f}<extra></extra>"
        ),
        row=1, col=1
    )
    
    # Highlight extreme regions
    for region_type, region in analysis.extreme_regions.items():
        color = 'rgba(255, 0, 0, 0.2)' if region_type == 'human' else 'rgba(0, 0, 255, 0.2)'
        fig.add_vrect(
            x0=region['start'],
            x1=region['end'],
            fillcolor=color,
            opacity=0.5,
            layer="below",
            line_width=0,
            row=1, col=1
        )
    
    # Calculate rolling GC content
    window = 100
    gc_content = np.array([
        compute_gc_content(analysis.sequence[i:i+window])
        for i in range(0, len(analysis.sequence) - window + 1, window)
    ])
    
    # GC content trace
    fig.add_trace(
        go.Scatter(
            x=np.arange(len(gc_content)) * window,
            y=gc_content,
            name="GC%",
            line=dict(color='rgba(44, 160, 44, 0.8)'),
            hovertemplate="Position: %{x}<br>GC%: %{y:.1f}%<extra></extra>"
        ),
        row=2, col=1
    )
    
    # Update layout
    fig.update_layout(
        height=800,
        title=dict(
            text=f"Genome Analysis Overview<br><sub>{analysis.header}</sub>",
            x=0.5
        ),
        showlegend=False,
        plot_bgcolor='white'
    )
    
    # Update axes
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    return fig

def create_kmer_importance_plot(analysis: SequenceAnalysis, top_k: int = 10) -> go.Figure:
    """Create interactive k-mer importance plot using Plotly"""
    # Get top k-mers by absolute SHAP value
    kmers = [''.join(p) for p in product("ACGT", repeat=4)]
    indices = np.argsort(np.abs(analysis.shap_values))[-top_k:]
    
    # Create DataFrame for plotting
    df = pd.DataFrame({
        'k-mer': [kmers[i] for i in indices],
        'SHAP': analysis.shap_values[indices]
    })
    
    # Create plot
    fig = px.bar(
        df,
        x='SHAP',
        y='k-mer',
        orientation='h',
        color='SHAP',
        color_continuous_scale='RdBu',
        title=f'Top {top_k} Most Influential k-mers'
    )
    
    # Update layout
    fig.update_layout(
        height=400,
        plot_bgcolor='white',
        yaxis_title='',
        xaxis_title='SHAP Value',
        coloraxis_showscale=False
    )
    
    return fig

def create_shap_distribution_plot(analysis: SequenceAnalysis) -> go.Figure:
    """Create SHAP distribution plot using Plotly"""
    fig = go.Figure()
    
    # Add histogram
    fig.add_trace(go.Histogram(
        x=analysis.shap_means,
        nbinsx=50,
        name='SHAP Values',
        marker_color='rgba(31, 119, 180, 0.6)'
    ))
    
    # Add vertical line at x=0
    fig.add_vline(
        x=0,
        line_dash="dash",
        line_color="red",
        annotation_text="Neutral",
        annotation_position="top"
    )
    
    # Update layout
    fig.update_layout(
        title='Distribution of SHAP Values',
        xaxis_title='SHAP Value',
        yaxis_title='Count',
        plot_bgcolor='white',
        height=400
    )
    
    return fig

###############################################################################
# 5. MAIN ANALYSIS
###############################################################################

def analyze_sequence(
    file_obj: str = None,
    fasta_text: str = "",
    window_size: int = 500,
    model_path: str = 'model.pt',
    scaler_path: str = 'scaler.pkl'
) -> SequenceAnalysis:
    """Main sequence analysis function"""
    # Handle input
    if fasta_text.strip():
        text = fasta_text.strip()
    elif file_obj is not None:
        with open(file_obj, 'r') as f:
            text = f.read()
    else:
        raise ValueError("No input provided")
    
    # Parse FASTA
    sequences = parse_fasta(text)
    if not sequences:
        raise ValueError("No valid FASTA sequences found")
    
    header, seq = sequences[0]
    
    # Load model and scaler
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    state_dict = torch.load(model_path, map_location=device)
    model = VirusClassifier(256).to(device)
    model.load_state_dict(state_dict)
    
    scaler = joblib.load(scaler_path)
    
    # Process sequence
    freq_vector = sequence_to_kmer_vector(seq)
    scaled_vector = scaler.transform(freq_vector.reshape(1, -1))
    x_tensor = torch.FloatTensor(scaled_vector).to(device)
    
    # Get SHAP values and classification
    shap_values, prob_human = calculate_shap_values(model, x_tensor)
    prob_nonhuman = 1.0 - prob_human
    
    # Get per-base SHAP scores
    shap_means = compute_positionwise_scores(seq, shap_values)
    
    # Find extreme regions
    extreme_regions = find_extreme_regions(shap_means, window_size)
    
    # Create analysis object
    return SequenceAnalysis(
        header=header,
        sequence=seq,
        length=len(seq),
        gc_content=compute_gc_content(seq),
        classification="Human" if prob_human > 0.5 else "Non-human",
        human_prob=prob_human,
        nonhuman_prob=prob_nonhuman,
        shap_values=shap_values,
        shap_means=shap_means,
        extreme_regions=extreme_regions
    )

###############################################################################
# 6. GRADIO INTERFACE
###############################################################################

def create_interface():
    """Create enhanced Gradio interface with improved layout and interactivity"""
    
    def process_sequence(
        file_obj: str,
        fasta_text: str,
        window_size: int,
        top_kmers: int
    ) -> Tuple[str, List[go.Figure]]:
        """Process sequence and return formatted results and plots"""
        try:
            # Run analysis
            analysis = analyze_sequence(
                file_obj=file_obj,
                fasta_text=fasta_text,
                window_size=window_size
            )
            
            # Format results text
            results = f"""
            ### Sequence Analysis Results
            
            **Basic Information**
            - Sequence: {analysis.header}
            - Length: {analysis.length:,} bases
            - GC Content: {analysis.gc_content:.1f}%
            
            **Classification**
            - Prediction: {analysis.classification}
            - Human Probability: {analysis.human_prob:.3f}
            - Non-human Probability: {analysis.nonhuman_prob:.3f}
            
            **Extreme Regions (window size: {window_size}bp)**
            Most Human-like Region:
            - Position: {analysis.extreme_regions['human']['start']:,} - {analysis.extreme_regions['human']['end']:,}
            - Average SHAP: {analysis.extreme_regions['human']['avg_shap']:.4f}
            
            Most Non-human-like Region:
            - Position: {analysis.extreme_regions['nonhuman']['start']:,} - {analysis.extreme_regions['nonhuman']['end']:,}
            - Average SHAP: {analysis.extreme_regions['nonhuman']['avg_shap']:.4f}
            """
            
            # Create plots
            genome_plot = create_genome_overview_plot(analysis)
            kmer_plot = create_kmer_importance_plot(analysis, top_kmers)
            dist_plot = create_shap_distribution_plot(analysis)
            
            return results, [genome_plot, kmer_plot, dist_plot], analysis
            
        except Exception as e:
            return f"Error: {str(e)}", [], None
    
    # Create theme and styling
    theme = gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="gray",
    ).set(
        body_text_color="gray-dark",
        background_fill_primary="*gray-50",
        block_shadow="*shadow-sm",
        block_background_fill="white",
    )
    
    # Build interface
    with gr.Blocks(theme=theme, css="""
        .container { margin: 0 auto; max-width: 1200px; padding: 20px; }
        .results { margin-top: 20px; }
        .plot-container { margin-top: 10px; }
    """) as interface:
        gr.Markdown("""
        # 🧬 Enhanced Virus Host Classifier
        
        This tool analyzes viral sequences to predict their host (human vs. non-human) and provides detailed visualizations 
        of the features influencing this classification. Upload or paste a FASTA sequence to begin.
        
        *Using advanced SHAP analysis and interactive visualizations for interpretable results.*
        """)
        
        # Input section
        with gr.Tab("Sequence Analysis"):
            with gr.Row():
                with gr.Column(scale=1):
                    file_input = gr.File(
                        label="Upload FASTA File",
                        file_types=[".fasta", ".fa", ".txt"],
                        type="filepath"
                    )
                    
                    text_input = gr.Textbox(
                        label="Or Paste FASTA Sequence",
                        placeholder=">sequence_name\nACGTACGT...",
                        lines=5
                    )
                    
                    with gr.Row():
                        window_size = gr.Slider(
                            minimum=100,
                            maximum=5000,
                            value=500,
                            step=100,
                            label="Window Size for Region Analysis"
                        )
                        
                        top_kmers = gr.Slider(
                            minimum=5,
                            maximum=30,
                            value=10,
                            step=1,
                            label="Number of Top k-mers to Display"
                        )
                    
                    analyze_btn = gr.Button(
                        "🔍 Analyze Sequence",
                        variant="primary"
                    )
                
                # Results section
                with gr.Column(scale=2):
                    results_text = gr.Markdown(
                        label="Analysis Results"
                    )
                    
                    # Plots
                    genome_plot = gr.Plot(
                        label="Genome Overview"
                    )
                    
                    with gr.Row():
                        kmer_plot = gr.Plot(
                            label="k-mer Importance"
                        )
                        dist_plot = gr.Plot(
                            label="SHAP Distribution"
                        )
        
        # Help tab
        with gr.Tab("Help & Information"):
            gr.Markdown("""
            ### 📖 How to Use This Tool
            
            1. **Input Your Sequence**
               - Upload a FASTA file or paste your sequence in FASTA format
               - The sequence should contain only ACGT bases (non-standard bases will be filtered)
            
            2. **Adjust Parameters**
               - Window Size: Controls the length of regions analyzed for extreme patterns
               - Top k-mers: Number of most influential sequence patterns to display
            
            3. **Interpret Results**
               - Classification: Predicted host (human vs. non-human)
               - Genome Overview: Interactive plot showing SHAP values and GC content
               - k-mer Importance: Most influential sequence patterns
               - SHAP Distribution: Overall distribution of feature importance
            
            ### 🎨 Visualization Guide
            
            - **SHAP Values**: 
              - Positive (red) = pushing toward human classification
              - Negative (blue) = pushing toward non-human classification
              - Zero (white) = neutral impact
            
            - **Extreme Regions**:
              - Highlighted in the genome overview plot
              - Red regions = most human-like
              - Blue regions = most non-human-like
            
            ### 🔬 Technical Details
            
            - The classifier uses k-mer frequencies (k=4) as features
            - SHAP values are calculated using an ablation-based approach
            - GC content is calculated using a sliding window
            """)
        
        # Connect components
        sequence_state = gr.State()
        
        analyze_btn.click(
            process_sequence,
            inputs=[
                file_input,
                text_input,
                window_size,
                top_kmers
            ],
            outputs=[
                results_text,
                [genome_plot, kmer_plot, dist_plot],
                sequence_state
            ]
        )
        
        return interface

###############################################################################
# 7. MAIN ENTRY POINT
###############################################################################

if __name__ == "__main__":
    iface = create_interface()
    iface.launch(
        share=True,
        server_name="0.0.0.0",
        show_error=True
    )
    #