import gradio as gr
import torch
import joblib
import numpy as np
from itertools import product
import torch.nn as nn
import shap
import matplotlib.pyplot as plt
import io
from PIL import Image  # Import PIL for image handling

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)
    
    def get_feature_importance(self, x):
        """Calculate feature importance using gradient-based method"""
        x.requires_grad_(True)
        output = self.network(x)
        importance = torch.zeros_like(x)
        
        for i in range(output.shape[1]):
            if x.grad is not None:
                x.grad.zero_()
            output[..., i].sum().backward(retain_graph=True)
            importance += torch.abs(x.grad)
            
        return importance

def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
    """Convert sequence to k-mer frequency vector"""
    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

def parse_fasta(text):
    sequences = []
    current_header = None
    current_sequence = []
    
    for line in text.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 predict(file_obj):
    if file_obj is None:
        return "Please upload a FASTA file", None
    
    try:
        if isinstance(file_obj, str):
            text = file_obj
        else:
            text = file_obj.decode('utf-8')
    except Exception as e:
        return f"Error reading file: {str(e)}", None

    k = 4
    kmers = [''.join(p) for p in product("ACGT", repeat=k)]
    kmer_dict = {km: i for i, km in enumerate(kmers)}
    
    try:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = VirusClassifier(256).to(device)
        state_dict = torch.load('model.pt', map_location=device)
        model.load_state_dict(state_dict)
        scaler = joblib.load('scaler.pkl')
        model.eval()
    except Exception as e:
        return f"Error loading model: {str(e)}", None

    results_text = ""
    plot_image = None
    
    try:
        sequences = parse_fasta(text)
        header, seq = sequences[0]
        
        raw_freq_vector = sequence_to_kmer_vector(seq)
        kmer_vector = scaler.transform(raw_freq_vector.reshape(1, -1))
        X_tensor = torch.FloatTensor(kmer_vector).to(device)
        
        with torch.no_grad():
            output = model(X_tensor)
            probs = torch.softmax(output, dim=1)
        
        importance = model.get_feature_importance(X_tensor)
        kmer_importance = importance[0].cpu().numpy()
        
        if np.max(np.abs(kmer_importance)) != 0:
            kmer_importance = kmer_importance / np.max(np.abs(kmer_importance)) * 0.002
        
        top_k = 10
        top_indices = np.argsort(np.abs(kmer_importance))[-top_k:][::-1]
        important_kmers = [
            {
                'kmer': list(kmer_dict.keys())[list(kmer_dict.values()).index(i)],
                'importance': float(kmer_importance[i]),
                'frequency': float(raw_freq_vector[i]),
                'scaled': float(kmer_vector[0][i])
            }
            for i in top_indices
        ]
        
        top_features = [item['kmer'] for item in important_kmers]
        top_values = [item['importance'] for item in important_kmers]
        
        others_mask = np.ones_like(kmer_importance, dtype=bool)
        others_mask[top_indices] = False
        others_sum = np.sum(kmer_importance[others_mask])
        
        top_features.append("Others")
        top_values.append(others_sum)
        
        # Set base_values and expected_value to 0.5 for the binary classification baseline
        explanation = shap.Explanation(
            values=np.array(top_values),
            base_values=0.5,  
            data=np.array([
                raw_freq_vector[kmer_dict[feat]] if feat != "Others" 
                else np.sum(raw_freq_vector[others_mask]) 
                for feat in top_features
            ]),
            feature_names=top_features
        )
        explanation.expected_value = 0.5
        
        fig = shap.plots._waterfall.waterfall_legacy(explanation, show=False)
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png')
        buf.seek(0)
        plot_image = Image.open(buf)

        pred_class = 1 if probs[0][1] > probs[0][0] else 0
        pred_label = 'human' if pred_class == 1 else 'non-human'
        
        results_text += f"""Sequence: {header}
Prediction: {pred_label}
Confidence: {float(max(probs[0])):0.4f}
Human probability: {float(probs[0][1]):0.4f}
Non-human probability: {float(probs[0][0]):0.4f}
Most influential k-mers (ranked by importance):"""
        
        for kmer in important_kmers:
            results_text += f"\n  {kmer['kmer']}: "
            results_text += f"impact={kmer['importance']:.4f}, "
            results_text += f"occurrence={kmer['frequency']*100:.2f}% of sequence "
            if kmer['scaled'] > 0:
                results_text += f"(appears {abs(kmer['scaled']):.2f}σ more than average)"
            else:
                results_text += f"(appears {abs(kmer['scaled']):.2f}σ less than average)"
        
    except Exception as e:
        return f"Error processing sequences: {str(e)}", None

    return results_text, plot_image

iface = gr.Interface(
    fn=predict,
    inputs=gr.File(label="Upload FASTA file", type="binary"),
    outputs=[gr.Textbox(label="Results"), gr.Image(label="SHAP Waterfall Plot")],
    title="Virus Host Classifier"
)

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