import streamlit as st
import numpy as np
import pandas as pd
import altair as alt
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# --------------------------
# Chargement du modèle rapide (FLAN-T5 Small)
# --------------------------
@st.cache_resource
def load_model():
    model_id = "google/flan-t5-small"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForSeq2SeqLM.from_pretrained(
        model_id,
        torch_dtype=torch.float32
    )
    return tokenizer, model

tokenizer, model = load_model()

# --------------------------
# Interface Streamlit
# --------------------------
st.set_page_config(layout="wide")
st.title("🌀 Spirale interactive + 🤖 Chatbot léger (FLAN-T5)")

# --------------------------
# Partie 1 : Spirale interactive
# --------------------------
with st.sidebar:
    st.header("🌀 Contrôle de la spirale")
    num_points = st.slider("Nombre de points", 1, 10000, 1100)
    num_turns = st.slider("Nombre de tours", 1, 300, 31)

indices = np.linspace(0, 1, num_points)
theta = 2 * np.pi * num_turns * indices
radius = indices

x = radius * np.cos(theta)
y = radius * np.sin(theta)

df = pd.DataFrame({
    "x": x,
    "y": y,
    "idx": indices,
    "rand": np.random.randn(num_points),
})

chart = alt.Chart(df, height=600, width=600).mark_point(filled=True).encode(
    x=alt.X("x", axis=None),
    y=alt.Y("y", axis=None),
    color=alt.Color("idx", legend=None, scale=alt.Scale(scheme='viridis')),
    size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
)

# Affichage de la spirale
st.subheader("🌀 Spirale générée")
st.altair_chart(chart)

# --------------------------
# Partie 2 : Chatbot FLAN-T5
# --------------------------
st.subheader("💬 Chat avec FLAN-T5 (Modèle rapide)")

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

user_input = st.text_input("Pose une question ou donne une consigne...", "")

if st.button("Envoyer") and user_input.strip():
    with st.spinner("Réflexion en cours..."):
        prompt = user_input.strip()
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
        output_ids = model.generate(input_ids, max_new_tokens=150)
        response = tokenizer.decode(output_ids[0], skip_special_tokens=True)

        st.session_state.chat_history.append(("👤", prompt))
        st.session_state.chat_history.append(("🤖", response))

# Affichage de l'historique du chat
for speaker, msg in st.session_state.chat_history:
    st.markdown(f"**{speaker}**: {msg}")