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():
    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:
        # Betöltés Excel fájl
        excel_data = pd.ExcelFile(uploaded_file)
        sheet_names = excel_data.sheet_names
        
        st.write("Munkalapok a fájlban:")
        st.write(sheet_names)

        # Minden munkalapot megvizsgál
        relevant_data = []
        for sheet in sheet_names:
            df = excel_data.parse(sheet)
            if "Név" in df.columns and "Elérhetőség" in df.columns:
                st.write(f"Adatok a '{sheet}' munkalapról:")
                st.dataframe(df)
                relevant_data.append(df)
        
        # Összesített adatok
        if relevant_data:
            combined_data = pd.concat(relevant_data, ignore_index=True)
            st.write("Kombinált adatok:")
            st.dataframe(combined_data)

            # Adatfeldolgozás
            combined_data["Elérhető napok"] = combined_data["Elérhetőség"].apply(lambda x: x.split(","))
            combined_data["Korlátozások"] = combined_data["Korlátozások"].apply(lambda x: x.split(",") if pd.notna(x) else [])
            foglalt_napok = {}
            beosztas = []

            for nap in range(1, 8):  # Példa: 7 napos beosztás
                nap_nev = f"Nap {nap}"
                for index, row in combined_data.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

            # Beosztás 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"
            )

        else:
            st.warning("Nem található releváns adat a munkalapokon.")

    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.")