import streamlit as st
import json
import pandas as pd
import numpy as np

st.set_page_config(page_title="PetSet", layout="centered")

st.markdown("<h1 style='text-align: center;'>🐾 PetSet – Your Dataset Pet</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center; font-size: 18px;'>Upload a LeRobot .json dataset and start caring for your pet!</p>", unsafe_allow_html=True)

uploaded_file = st.file_uploader("📂 Upload your LeRobotDataset (.json)", type=["json"])

if uploaded_file:
    data = json.load(uploaded_file)
    episodes = data.get("episodes", [])

    if not episodes:
        st.error("No episodes found in this file.")
    else:
        df = pd.DataFrame(episodes)
        st.success(f"{len(episodes)} episodes loaded successfully!")

        # Calcular atributos do pet
        success_rate = df["success"].mean()
        duration_std = df["duration"].std()
        health = 100 - df['success'].value_counts().get(False, 0) * 10
        energy = success_rate * 100
        attention = max(0, 100 - duration_std * 10)

        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("❤️ Health", f"{int(health)}%")
        with col2:
            st.metric("⚡ Energy", f"{int(energy)}%")
        with col3:
            st.metric("🧠 Attention", f"{int(attention)}%")

        st.markdown("---")
        st.image("https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExZ2EyZWkyd2ZzMmFmd3FuNm1jY2Fjdm42a2p6dDh4cGc1dzNsb25uMSZlcD12MV9naWZzX3NlYXJjaCZjdD1n/VbnUQpnihPSIgIXuZv/giphy.gif", width=200, caption="Your Pet")

        st.markdown("### 🛠 What would you like to do?")
        col_a, col_b, col_c = st.columns(3)
        with col_a:
            feed = st.button("🍎 Feed Data")
        with col_b:
            heal = st.button("🩺 Heal Pet")
        with col_c:
            check = st.button("🧪 Check Quality")

        if heal:
            original_len = len(data["episodes"])
            data["episodes"] = [ep for ep in episodes if ep.get("success", True)]
            removed = original_len - len(data["episodes"])
            st.success(f"Removed {removed} failed episodes.")

        if check:
            problems = df['success'].value_counts().get(False, 0)
            st.markdown(f"<div style='padding: 15px; background-color: #fff3cd; border-radius: 10px;'><strong>⚠️ Found {problems} corrupted episodes!</strong></div>", unsafe_allow_html=True)

        st.markdown("### 📦 Dataset Preview")
        st.dataframe(df[["id", "success", "duration", "avg_accel"]])

        st.markdown("---")
        if st.button("💾 Download Cleaned Dataset"):
            cleaned_json = json.dumps(data, indent=2)
            st.download_button("Download JSON", cleaned_json, file_name="cleaned_dataset.json")

else:
    st.info("👆 Upload a dataset file to begin.")