app.py

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 io
from PIL import Image

###############################################################################
# 1. MODEL DEFINITION
###############################################################################

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. FASTA PARSING & K-MER FEATURE ENGINEERING
###############################################################################

def parse_fasta(text):
    """Parse FASTA formatted text into a list of (header, sequence)."""
    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:
            current_sequence.append(line.upper())
    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 a sequence to a k-mer frequency vector for classification."""
    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)
    
    for i in range(len(sequence) - k + 1):
        kmer = sequence[i:i+k]
        if kmer in kmer_dict:
            vec[kmer_dict[kmer]] += 1

    total_kmers = len(sequence) - k + 1
    if total_kmers > 0:
        vec = vec / total_kmers

    return vec

###############################################################################
# 3. SHAP-VALUE (ABLATION) CALCULATION
###############################################################################

def calculate_shap_values(model, x_tensor):
    """
    Calculate SHAP values using a simple ablation approach.
    Returns shap_values, prob_human
    """
    model.eval()
    with torch.no_grad():
        # Baseline
        baseline_output = model(x_tensor)
        baseline_probs = torch.softmax(baseline_output, dim=1)
        baseline_prob = baseline_probs[0, 1].item()  # Probability of 'human' class
        
        # Zeroing each feature to measure impact
        shap_values = []
        x_zeroed = x_tensor.clone()
        for i in range(x_tensor.shape[1]):
            original_val = x_zeroed[0, i].item()
            x_zeroed[0, i] = 0.0
            output = model(x_zeroed)
            probs = torch.softmax(output, dim=1)
            prob = probs[0, 1].item()
            impact = baseline_prob - prob
            shap_values.append(impact)
            x_zeroed[0, i] = original_val  # restore
    return np.array(shap_values), baseline_prob

###############################################################################
# 4. PER-BASE SHAP AGGREGATION
###############################################################################

def compute_positionwise_scores(sequence, shap_values, k=4):
    """
    Returns an array of per-base SHAP contributions by averaging
    the k-mer SHAP values of all k-mers covering that base.
    """
    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)
    
    for i in range(seq_len - k + 1):
        kmer = sequence[i:i+k]
        if kmer in kmer_dict:
            val = shap_values[kmer_dict[kmer]]
            shap_sums[i : i + k] += val
            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

###############################################################################
# 5. PLOTTING / UTILITIES
###############################################################################

def fig_to_image(fig):
    """Convert a Matplotlib figure to a PIL Image for Gradio."""
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=150)
    buf.seek(0)
    img = Image.open(buf)
    plt.close(fig)
    return img

def plot_linear_heatmap(shap_means, title="Per-base SHAP Heatmap", start=None, end=None):
    """
    Plots a 1D heatmap of per-base SHAP contributions.
    Negative = push toward Non-Human, Positive = push toward Human.
    Optionally can show only a subrange (start:end).
    """
    if start is not None and end is not None:
        shap_means = shap_means[start:end]
        subtitle = f" (positions {start}-{end})"
    else:
        subtitle = ""
    
    heatmap_data = shap_means.reshape(1, -1)  # shape (1, region_length)
    
    fig, ax = plt.subplots(figsize=(12, 2))
    cax = ax.imshow(heatmap_data, aspect='auto', cmap='RdBu_r')
    cbar = plt.colorbar(cax, orientation='horizontal', pad=0.2)
    cbar.set_label('SHAP Contribution')

    ax.set_yticks([])
    ax.set_xlabel('Position in Sequence')
    ax.set_title(f"{title}{subtitle}")
    plt.tight_layout()
    return fig

def create_importance_bar_plot(shap_values, kmers, top_k=10):
    """Create a bar plot of the most important k-mers."""
    plt.rcParams.update({'font.size': 10})
    fig = plt.figure(figsize=(10, 5))
    
    # Sort by absolute importance
    indices = np.argsort(np.abs(shap_values))[-top_k:]
    values = shap_values[indices]
    features = [kmers[i] for i in indices]
    
    colors = ['#ff9999' if v > 0 else '#99ccff' for v in values]
    
    plt.barh(range(len(values)), values, color=colors)
    plt.yticks(range(len(values)), features)
    plt.xlabel('SHAP Value (impact on model output)')
    plt.title(f'Top {top_k} Most Influential k-mers')
    plt.gca().invert_yaxis()
    return fig

def compute_gc_content(sequence):
    """Compute %GC in the sequence (A, C, G, T)."""
    if not sequence:
        return 0
    gc_count = sequence.count('G') + sequence.count('C')
    return (gc_count / len(sequence)) * 100.0

###############################################################################
# 6. MAIN ANALYSIS STEP (Gradio Step 1)
###############################################################################

def analyze_sequence(file_obj, top_kmers=10, fasta_text=""):
    """Analyzes the entire genome, returning classification and a heatmap."""
    # Handle input
    if fasta_text.strip():
        text = fasta_text.strip()
    elif file_obj is not None:
        try:
            with open(file_obj, 'r') as f:
                text = f.read()
        except Exception as e:
            return (f"Error reading file: {str(e)}", None, None, None, None)
    else:
        return ("Please provide a FASTA sequence.", None, None, None, None)

    # Parse FASTA
    sequences = parse_fasta(text)
    if not sequences:
        return ("No valid FASTA sequences found.", None, None, None, None)
    
    header, seq = sequences[0]

    # Load model and scaler
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        model = VirusClassifier(256).to(device)
        model.load_state_dict(torch.load('model.pt', map_location=device))
        scaler = joblib.load('scaler.pkl')
    except Exception as e:
        return (f"Error loading model: {str(e)}", None, None, None, None)

    # Vectorize + scale
    freq_vector = sequence_to_kmer_vector(seq)
    scaled_vector = scaler.transform(freq_vector.reshape(1, -1))
    x_tensor = torch.FloatTensor(scaled_vector).to(device)

    # SHAP + classification
    shap_values, prob_human = calculate_shap_values(model, x_tensor)
    prob_nonhuman = 1.0 - prob_human
    
    classification = "Human" if prob_human > 0.5 else "Non-human"
    confidence = max(prob_human, prob_nonhuman)

    # Build results text
    results_text = (
        f"Sequence: {header}\n"
        f"Length: {len(seq):,} bases\n"
        f"Classification: {classification}\n"
        f"Confidence: {confidence:.3f}\n"
        f"(Human Probability: {prob_human:.3f}, Non-human Probability: {prob_nonhuman:.3f})"
    )

    # K-mer importance plot
    kmers = [''.join(p) for p in product("ACGT", repeat=4)]
    bar_fig = create_importance_bar_plot(shap_values, kmers, top_kmers)
    bar_img = fig_to_image(bar_fig)

    # Per-base SHAP for entire genome
    shap_means = compute_positionwise_scores(seq, shap_values, k=4)
    heatmap_fig = plot_linear_heatmap(shap_means, title="Genome-wide SHAP")
    heatmap_img = fig_to_image(heatmap_fig)

    # Return:
    # 1) results text
    # 2) k-mer bar image
    # 3) full-genome heatmap
    # 4) the "state" we need for step 2: (sequence, shap_means)
    #    We'll store these in a dictionary so we can pass it around in Gradio.
    state_dict = {
        "seq": seq,
        "shap_means": shap_means
    }

    return (results_text, bar_img, heatmap_img, state_dict, header)

###############################################################################
# 7. SUBREGION ANALYSIS (Gradio Step 2)
###############################################################################

def analyze_subregion(state, header, region_start, region_end):
    """
    Takes stored data from step 1 and a user-chosen region.
    Returns a subregion heatmap and some stats (like GC content, average SHAP).
    """
    if not state or "seq" not in state or "shap_means" not in state:
        return ("No sequence data found. Please run Step 1 first.", None)
    
    seq = state["seq"]
    shap_means = state["shap_means"]

    # Validate bounds
    region_start = max(0, min(region_start, len(seq)))
    region_end = max(0, min(region_end, len(seq)))
    if region_end <= region_start:
        return ("Invalid region range. End must be > Start.", None)

    # Subsequence
    region_seq = seq[region_start:region_end]
    region_shap = shap_means[region_start:region_end]

    # Some stats
    gc_percent = compute_gc_content(region_seq)
    avg_shap = float(np.mean(region_shap))
    
    region_info = (
        f"Analyzing subregion of {header} from {region_start} to {region_end}\n"
        f"Region length: {len(region_seq)} bases\n"
        f"GC content: {gc_percent:.2f}%\n"
        f"Average SHAP in region: {avg_shap:.4f} "
        f"({'toward human' if avg_shap > 0 else 'toward non-human' if avg_shap < 0 else 'neutral'})"
    )

    # Plot region as small heatmap
    fig = plot_linear_heatmap(shap_means, 
                              title="Subregion SHAP", 
                              start=region_start, 
                              end=region_end)
    heatmap_img = fig_to_image(fig)

    return (region_info, heatmap_img)


###############################################################################
# 8. BUILD GRADIO INTERFACE
###############################################################################

css = """
.gradio-container {
    font-family: 'IBM Plex Sans', sans-serif;
}
"""

with gr.Blocks(css=css) as iface:
    gr.Markdown("""
    # Virus Host Classifier (with Interactive Region Viewer)
    **Step 1**: Predict overall viral sequence origin (human vs non-human)  
    **Step 2**: Explore subregions to see local SHAP signals and GC content
    """)
    
    with gr.Tab("1) Full-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
                )
                top_k = 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")
                
            with gr.Column(scale=2):
                results_box = gr.Textbox(
                    label="Classification Results", lines=7, interactive=False
                )
                kmer_img = gr.Image(label="Top k-mer SHAP")
                genome_img = gr.Image(label="Genome-wide SHAP Heatmap")
        
        # Hidden states that store data for step 2
        # "state" will hold (sequence, shap_means).
        # "header" is optional meta info
        seq_state = gr.State()
        header_state = gr.State()

        # The "analyze_sequence" function returns 5 values, which we map here:
        analyze_btn.click(
            analyze_sequence,
            inputs=[file_input, top_k, text_input],
            outputs=[results_box, kmer_img, genome_img, seq_state, header_state]
        )
    
    with gr.Tab("2) Subregion Exploration"):
        gr.Markdown("""
        Select start/end positions to view local SHAP signals.
        """)
        with gr.Row():
            region_start = gr.Number(label="Region Start", value=0)
            region_end = gr.Number(label="Region End", value=500)
            region_btn = gr.Button("Analyze Subregion")
        
        subregion_info = gr.Textbox(
            label="Subregion Analysis",
            lines=4, 
            interactive=False
        )
        subregion_img = gr.Image(label="Subregion SHAP Heatmap")

        region_btn.click(
            analyze_subregion,
            inputs=[seq_state, header_state, region_start, region_end],
            outputs=[subregion_info, subregion_img]
        )
    
    gr.Markdown("""
    ### What does this interface provide?
    1. **Overall Classification** (human vs non-human), using a learned model on k-mer frequencies.
    2. **SHAP Analysis** (ablation-based) to see which k-mer features push classification toward or away from "human".
    3. **Genome-Wide SHAP Heatmap**: Each base's average SHAP across overlapping k-mers.
    4. **Subregion Exploration**:  
       - View SHAP signals in a user-chosen region.  
       - Calculate local GC content, average SHAP, etc.  

    ### Tips
    - For very large sequences (e.g., >100k bases), the full heatmap might be large; consider downsampling if needed.
    - Adjust *Region Start* and *End* to explore different parts of the genome.
    """)

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