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 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. FIND EXTREME SHAP REGIONS
###############################################################################

def find_extreme_subregion(shap_means, window_size=500, mode="max"):
    """
    Finds the subregion of length `window_size` that has the maximum 
    (mode="max") or minimum (mode="min") average SHAP.
    Returns (best_start, best_end, best_avg).
    """
    n = len(shap_means)
    if n == 0:
        return (0, 0, 0.0)
    if window_size >= n:
        # entire sequence
        avg_val = float(np.mean(shap_means))
        return (0, n, avg_val)
    
    # We'll build csum of length n+1
    csum = np.zeros(n + 1, dtype=np.float32)
    csum[1:] = np.cumsum(shap_means)

    best_start = 0
    best_sum = csum[window_size] - csum[0]
    best_avg = best_sum / window_size
    
    for start in range(1, n - window_size + 1):
        wsum = csum[start + window_size] - csum[start]
        wavg = wsum / window_size
        if mode == "max":
            if wavg > best_avg:
                best_avg = wavg
                best_start = start
        else:  # mode == "min"
            if wavg < best_avg:
                best_avg = wavg
                best_start = start

    return (best_start, best_start + window_size, float(best_avg))

###############################################################################
# 6. 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 get_zero_centered_cmap():
    """
    Creates a custom diverging colormap that is:
    - Blue for negative
    - White for zero
    - Red for positive
    """
    colors = [
        (0.0, 'blue'),   # negative
        (0.5, 'white'),  # zero
        (1.0, 'red')     # positive
    ]
    cmap = mcolors.LinearSegmentedColormap.from_list("blue_white_red", colors)
    return cmap

def plot_linear_heatmap(shap_means, title="Per-base SHAP Heatmap", start=None, end=None):
    """
    Plots a 1D heatmap of per-base SHAP contributions with a custom colormap:
    - Negative = blue
    - 0 = white
    - Positive = red
    """
    if start is not None and end is not None:
        local_shap = shap_means[start:end]
        subtitle = f" (positions {start}-{end})"
    else:
        local_shap = shap_means
        subtitle = ""
        
    if len(local_shap) == 0:
        local_shap = np.array([0.0])
        
    # Build 2D array for imshow
    heatmap_data = local_shap.reshape(1, -1)
    
    # Force symmetrical range
    min_val = np.min(local_shap)
    max_val = np.max(local_shap)
    extent = max(abs(min_val), abs(max_val))
    
    # Create custom colormap
    custom_cmap = get_zero_centered_cmap()
    
    # Create figure with adjusted height ratio
    fig, ax = plt.subplots(figsize=(12, 1.8))  # Reduced height
    
    # Plot heatmap
    cax = ax.imshow(
        heatmap_data,
        aspect='auto',
        cmap=custom_cmap,
        vmin=-extent,
        vmax=+extent
    )
    
    # Configure colorbar with more subtle positioning
    cbar = plt.colorbar(
        cax,
        orientation='horizontal',
        pad=0.25,  # Reduced padding
        aspect=40,  # Make colorbar thinner
        shrink=0.8  # Make colorbar shorter than plot width
    )
    
    # Style the colorbar
    cbar.ax.tick_params(labelsize=8)  # Smaller tick labels
    cbar.set_label(
        'SHAP Contribution',
        fontsize=9,
        labelpad=5
    )
    
    # Configure main plot
    ax.set_yticks([])
    ax.set_xlabel('Position in Sequence', fontsize=10)
    ax.set_title(f"{title}{subtitle}", pad=10)
    
    # Fine-tune layout
    plt.subplots_adjust(
        bottom=0.25,  # Reduced bottom margin
        left=0.05,    # Tighter left margin
        right=0.95    # Tighter right margin
    )
    
    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]
    
    # negative -> blue, positive -> red
    colors = ['#99ccff' if v < 0 else '#ff9999' 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()
    plt.tight_layout()
    return fig

def plot_shap_histogram(shap_array, title="SHAP Distribution in Region"):
    """
    Simple histogram of SHAP values in the subregion.
    """
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(shap_array, bins=30, color='gray', edgecolor='black')
    ax.axvline(0, color='red', linestyle='--', label='0.0')
    ax.set_xlabel("SHAP Value")
    ax.set_ylabel("Count")
    ax.set_title(title)
    ax.legend()
    plt.tight_layout()
    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

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

def analyze_sequence(file_obj, top_kmers=10, fasta_text="", window_size=500):
    """
    Analyzes the entire genome, returning classification, full-genome heatmap,
    top k-mer bar plot, and identifies subregions with strongest positive/negative push.
    """
    # 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:
        # Use weights_only=True for safer loading
        state_dict = torch.load('model.pt', map_location=device, weights_only=True)
        model = VirusClassifier(256).to(device)
        model.load_state_dict(state_dict)
        
        scaler = joblib.load('scaler.pkl')
    except Exception as e:
        return (f"Error loading model/scaler: {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)

    # Per-base SHAP
    shap_means = compute_positionwise_scores(seq, shap_values, k=4)

    # Find the most "human-pushing" region
    (max_start, max_end, max_avg) = find_extreme_subregion(shap_means, window_size, mode="max")
    # Find the most "non-human–pushing" region
    (min_start, min_end, min_avg) = find_extreme_subregion(shap_means, window_size, mode="min")

    # 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})\n\n"
        f"---\n"
        f"**Most Human-Pushing {window_size}-bp Subregion**:\n"
        f"Start: {max_start}, End: {max_end}, Avg SHAP: {max_avg:.4f}\n\n"
        f"**Most Non-Human–Pushing {window_size}-bp Subregion**:\n"
        f"Start: {min_start}, End: {min_end}, Avg SHAP: {min_avg:.4f}"
    )

    # 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)

    # Full-genome SHAP heatmap
    heatmap_fig = plot_linear_heatmap(shap_means, title="Genome-wide SHAP")
    heatmap_img = fig_to_image(heatmap_fig)

    # Store data for subregion analysis
    state_dict_out = {
        "seq": seq,
        "shap_means": shap_means
    }

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

###############################################################################
# 8. 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, histogram, and some stats (GC, 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, None)
    
    seq = state["seq"]
    shap_means = state["shap_means"]

    # Validate bounds
    region_start = int(region_start)
    region_end = int(region_end)

    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, 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))

    # Fraction pushing toward human vs. non-human
    positive_fraction = np.mean(region_shap > 0)
    negative_fraction = np.mean(region_shap < 0)

    # Simple logic-based interpretation
    if avg_shap > 0.05:
        region_classification = "Likely pushing toward human"
    elif avg_shap < -0.05:
        region_classification = "Likely pushing toward non-human"
    else:
        region_classification = "Near neutral (no strong push)"

    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}\n"
        f"Fraction with SHAP > 0 (toward human): {positive_fraction:.2f}\n"
        f"Fraction with SHAP < 0 (toward non-human): {negative_fraction:.2f}\n"
        f"Subregion interpretation: {region_classification}\n"
    )

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

    # Plot histogram of SHAP in region
    hist_fig = plot_shap_histogram(region_shap, title="SHAP Distribution in Subregion")
    hist_img = fig_to_image(hist_fig)

    return (region_info, heatmap_img, hist_img)


###############################################################################
# 9. 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
    **Step 1**: Predict overall viral sequence origin (human vs non-human) and identify extreme regions.  
    **Step 2**: Explore subregions to see local SHAP signals, distribution, GC content, etc.
    
    **Color Scale**: Negative SHAP = Blue, Zero = White, Positive = Red.
    """)
    
    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"
                )
                win_size = gr.Slider(
                    minimum=100,
                    maximum=5000,
                    value=500,
                    step=100,
                    label="Window size for 'most pushing' subregions"
                )
                analyze_btn = gr.Button("Analyze Sequence", variant="primary")
                
            with gr.Column(scale=2):
                results_box = gr.Textbox(
                    label="Classification Results", lines=12, interactive=False
                )
                kmer_img = gr.Image(label="Top k-mer SHAP")
                genome_img = gr.Image(label="Genome-wide SHAP Heatmap (Blue=neg, White=0, Red=pos)")
        
        seq_state = gr.State()
        header_state = gr.State()

        # analyze_sequence(...) returns 5 items
        analyze_btn.click(
            analyze_sequence,
            inputs=[file_input, top_k, text_input, win_size],
            outputs=[results_box, kmer_img, genome_img, seq_state, header_state]
        )
    
    with gr.Tab("2) Subregion Exploration"):
        gr.Markdown("""
        **Subregion Analysis**  
        Select start/end positions to view local SHAP signals, distribution, and GC content.  
        The heatmap also uses the same Blue-White-Red scale.  
        """)
        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=7, 
            interactive=False
        )
        with gr.Row():
            subregion_img = gr.Image(label="Subregion SHAP Heatmap (B-W-R)")
            subregion_hist_img = gr.Image(label="SHAP Distribution (Histogram)")
        
        region_btn.click(
            analyze_subregion,
            inputs=[seq_state, header_state, region_start, region_end],
            outputs=[subregion_info, subregion_img, subregion_hist_img]
        )
    
    gr.Markdown("""
    ### Interface Features
    - **Overall Classification** (human vs non-human) using k-mer frequencies.
    - **SHAP Analysis** to see which k-mers push classification toward or away from human.
    - **White-Centered SHAP Gradient**: 
      - Negative (blue), 0 (white), Positive (red), with symmetrical color range around 0. 
    - **Identify Subregions** with the strongest push for human or non-human.
    - **Subregion Exploration**: 
      - Local SHAP heatmap & histogram 
      - GC content 
      - Fraction of positions pushing human vs. non-human 
      - Simple logic-based classification
    """)

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