import streamlit as st
import torch
import pandas as pd
from utils import load_model_and_tokenizer, preprocess_input, predict_property
from attention_plot import plot_attention_heatmap

st.set_page_config(page_title="TransPolymer Predictor", layout="centered")
st.title("🔬 TransPolymer Property Prediction")
st.markdown("Masukkan struktur polimer (SMILES) dan maklumat deskriptor.")

tokenizer, model = load_model_and_tokenizer()

# Input form
with st.form("polymer_form"):
    smiles = st.text_input("📄 Polymer SMILES", value="*CCO*")
    temp = st.number_input("🌡️ Suhu (°C)", value=25.0)
    mw = st.number_input("⚖️ Berat Molekul (kDa)", value=100.0)
    submitted = st.form_submit_button("Jalankan Ramalan")

if submitted:
    input_seq = preprocess_input(smiles, temp, mw)
    pred = predict_property(input_seq, model, tokenizer)
    st.success(f"📊 Ramalan Sifat Polimer: **{pred:.4f}**")
    st.markdown("### 🎯 Visualisasi Attention (Last Layer)")
    fig = plot_attention_heatmap(model, tokenizer, input_seq)
    st.pyplot(fig)

# Batch prediction
st.markdown("---")
st.header("📁 Batch Prediction (CSV)")
uploaded_file = st.file_uploader("Muat naik CSV dengan lajur: `smiles`, `temp`, `mw`", type="csv")

if uploaded_file:
    df = pd.read_csv(uploaded_file)
    results = []
    for i, row in df.iterrows():
        seq = preprocess_input(row["smiles"], row["temp"], row["mw"])
        pred = predict_property(seq, model, tokenizer)
        results.append(pred)
    df["prediction"] = results
    st.dataframe(df)
    st.download_button("📥 Muat turun hasil", df.to_csv(index=False), "predictions.csv")