import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
@gr.on_startup()
def load_model():
    global tokenizer, model
    tokenizer = AutoTokenizer.from_pretrained("web3se/SmartBERT-v3")
    model = AutoModelForSequenceClassification.from_pretrained("web3se/SmartBERT-v3")
    model.eval()
    return model, tokenizer

# Function to analyze smart contract
def analyze_contract(contract_code):
    if not contract_code:
        return {"error": "No contract code provided"}
    
    try:
        # Tokenize input
        inputs = tokenizer(contract_code, return_tensors="pt", truncation=True, max_length=512)
        
        # Perform inference
        with torch.no_grad():
            outputs = model(**inputs)
        
        # Process outputs
        logits = outputs.logits
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        
        # Get predicted class and confidence
        predicted_class = torch.argmax(probabilities, dim=-1).item()
        confidence = probabilities[0, predicted_class].item()
        
        # Map class index to label
        labels = model.config.id2label
        predicted_label = labels[predicted_class]
        
        # Format all class probabilities for display
        all_probs = {labels[i]: f"{prob.item()*100:.2f}%" for i, prob in enumerate(probabilities[0])}
        sorted_probs = dict(sorted(all_probs.items(), key=lambda item: float(item[1].rstrip('%')), reverse=True))
        
        # Format the result as markdown for better display
        result_md = f"## Analysis Results\n\n"
        result_md += f"**Prediction:** {predicted_label}\n\n"
        result_md += f"**Confidence:** {confidence*100:.2f}%\n\n"
        result_md += "### All Class Probabilities:\n\n"
        
        for label, prob in sorted_probs.items():
            result_md += f"- {label}: {prob}\n"
            
        return result_md
        
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
title = "Smart Contract Analyzer (SmartBERT-v3)"
description = """
This app uses web3se/SmartBERT-v3 model to analyze smart contracts.
Simply paste your smart contract code in the text area below and click "Analyze".
"""

iface = gr.Interface(
    fn=analyze_contract,
    inputs=gr.Textbox(lines=15, placeholder="Paste your smart contract code here...", label="Smart Contract Code"),
    outputs=gr.Markdown(label="Analysis Results"),
    title=title,
    description=description,
    examples=[
        # You can add example smart contracts here
        ["pragma solidity ^0.8.0;\n\ncontract SimpleStorage {\n    uint256 private value;\n    \n    function set(uint256 _value) public {\n        value = _value;\n    }\n    \n    function get() public view returns (uint256) {\n        return value;\n    }\n}"]
    ],
    allow_flagging="never"
)

# Launch the app
iface.launch()