import streamlit as st
import pandas as pd
from transformers import pipeline

# Oldal konfiguráció
st.set_page_config(page_title="Orvosi Ügyeleti Beosztás Generáló", layout="wide")

st.title("Orvosi Ügyeleti Beosztás Generáló")
st.write("Töltsd fel az orvosi adatokat tartalmazó Excel fájlt, és generálj beosztást az ügyeletekhez!")

# Hugging Face nyelvi modell betöltése
@st.cache_resource
def load_model():
    # Kompatibilis modell betöltése (pl. gpt2)
    return pipeline("text-generation", model="gpt2")

generator = load_model()

# Fájl feltöltése
uploaded_file = st.file_uploader("Tölts fel egy Excel fájlt", type=["xlsx"])

if uploaded_file:
    try:
        # Excel beolvasása
        df = pd.read_excel(uploaded_file)
        st.write("Feltöltött adatok:")
        st.dataframe(df)

        # Adatok előkészítése
        df["Elérhető napok"] = df["Elérhetőség"].apply(lambda x: x.split(","))
        df["Korlátozások"] = df["Korlátozások"].apply(lambda x: x.split(",") if pd.notna(x) else [])
        foglalt_napok = {}
        beosztas = []

        # Beosztás generálása
        for nap in range(1, 8):  # Példa: 7 napos beosztás
            nap_nev = f"Nap {nap}"
            for index, row in df.iterrows():
                if nap_nev in row["Elérhető napok"]:
                    korlatozott = any(
                        foglalt_napok.get(nap_nev) == szemely
                        for szemely in row["Korlátozások"]
                    )
                    if not korlatozott:
                        foglalt_napok[nap_nev] = row["Név"]
                        # Szöveg generálása
                        prompt = f"{row['Név']} ügyel {nap_nev}-n, mert "
                        indoklas = generator(prompt, max_length=50, num_return_sequences=1)[0]["generated_text"]
                        beosztas.append({
                            "Nap": nap_nev,
                            "Orvos": row["Név"],
                            "Indoklás": indoklas
                        })
                        break

        # Eredmény megjelenítése
        beosztas_df = pd.DataFrame(beosztas)
        st.write("Generált Ügyeleti Beosztás:")
        st.dataframe(beosztas_df)

        # Exportálás Excelbe
        @st.cache_data
        def convert_to_excel(data):
            import io
            output = io.BytesIO()
            with pd.ExcelWriter(output, engine="openpyxl") as writer:
                data.to_excel(writer, index=False)
            processed_data = output.getvalue()
            return processed_data

        excel_data = convert_to_excel(beosztas_df)
        st.download_button(
            label="Beosztás letöltése Excelben",
            data=excel_data,
            file_name="ugyeleti_beosztas.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        )

    except Exception as e:
        st.error(f"Hiba történt a fájl feldolgozása során: {e}")
else:
    st.info("Tölts fel egy fájlt a kezdéshez.")