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

# 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üggvény a dátumok felismerésére és konverziójára
def parse_date(value):
    try:
        if any(char.isdigit() for char in value):
            return datetime.strptime(value, "%d %B %Y")
        parsed_date = datetime.strptime(value, "%d %B")
        return parsed_date.replace(year=2025)
    except (ValueError, TypeError):
        return None

# Adatok tisztítása és egységesítése
def clean_data(df):
    st.write("Eredeti oszlopnevek:", df.columns.tolist())
    df = df.dropna(how='all')  # Üres sorok eltávolítása
    df.columns = df.columns.str.strip().str.lower()
    st.write("Tisztított oszlopnevek:", df.columns.tolist())
    df = df.rename(columns={
        "orvos": "Név", 
        "név": "Név", 
        "elérhetőség": "Elérhetőség", 
        "korlátozások": "Korlátozások"
    })
    if "Név" not in df.columns:
        st.warning("A munkalap nem tartalmaz 'Név' oszlopot, ellenőrizze az adatok helyességét!")
    df = df.dropna(subset=["Név"], errors='ignore')  # Üres név mezők eltávolítása
    return df

# Adatok Excel exportálása
@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)
    return output.getvalue()

# 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)
            st.write(f"Feldolgozás alatt: {sheet}")
            st.write("Adatok előnézete:", df.head())
            df = clean_data(df)  # Adatok tisztítása és egységesítése

            # Ellenőrzés, hogy tartalmaz-e releváns oszlopokat
            if any(col in df.columns for col in ["Név", "név", "orvos"]) and "elérhetőség" in df.columns:
                st.write(f"Releváns adat található a '{sheet}' munkalapon.")
                relevant_data.append(df)
            else:
                st.warning(f"A '{sheet}' munkalap nem tartalmaz releváns adatokat.")

        # Összesített adatok
        if relevant_data:
            combined_data = pd.concat(relevant_data, ignore_index=True)
            st.write("Kombinált adatok:")
            st.dataframe(combined_data)

            # Dátumok feldolgozása az elérhetőség oszlopban
            combined_data["Elérhető napok"] = combined_data["Elérhetőség"].apply(
                lambda x: [parse_date(day.strip()) for day in str(x).split(",")] if pd.notna(x) else []
            )
            combined_data["Korlátozások"] = combined_data["Korlátozások"].apply(
                lambda x: x.split(",") if isinstance(x, str) else []
            )

            # Felhasználói input: dátumok kiválasztása
            available_days = st.multiselect(
                "Add meg az elérhetőségi dátumokat:",
                options=[day.strftime("%Y-%m-%d") for day_list in combined_data["Elérhető napok"] for day in day_list if day]
            )

            if available_days:
                foglalt_napok = {}
                beosztas = []

                for nap in available_days:
                    for index, row in combined_data.iterrows():
                        if nap in [day.strftime("%Y-%m-%d") for day in row["Elérhető napok"] if day]:
                            korlatozott = any(
                                foglalt_napok.get(nap) == szemely
                                for szemely in row["Korlátozások"]
                            )
                            if not korlatozott:
                                foglalt_napok[nap] = row["Név"]
                                prompt = f"{row['Név']} ügyel {nap}-n, mert "
                                indoklas = generator(prompt, max_length=50, num_return_sequences=1)[0]["generated_text"]
                                beosztas.append({
                                    "Nap": nap,
                                    "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
                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("Kérlek, válaszd ki az elérhetőségi dátumokat.")
        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.")