import streamlit as st
import pandas as pd
from datetime import timedelta

# Function to parse the uploaded file and return a DataFrame
def parse_roster(file):
    df = pd.read_excel(file)
    return df

# Function to generate the next one-year roster
def generate_yearly_roster(df):
    days_in_year = 365
    new_roster = []
    nurses = df['NAMES'].unique()

    for nurse in nurses:
        current_day = 0
        nurse_schedule = []

        while current_day < days_in_year:
            if nurse == "ADOMAKO SIMON" or nurse == "MAHAMA RAHMAT":
                shift_pattern = ['M', 'M', 'M', 'M', 'M', 'X', 'X']
            else:
                shift_pattern = ['M', 'A', 'N', 'N', 'N', 'N', 'N', 'X', 'X', 'M', 'A', 'A', 'M', 'X', 'X']

            for shift in shift_pattern:
                if current_day >= days_in_year:
                    break
                nurse_schedule.append(shift)
                current_day += 1

        new_roster.append(nurse_schedule)

    return pd.DataFrame(new_roster, index=nurses).T

# Streamlit app layout
st.title("Nurse Duty Roster Generator")

st.markdown("""
This app generates a one-year duty roster for nurses based on the current month's schedule.
* **Upload**: Upload the current month's roster in Excel format.
* **Generate**: Generate the duty roster for the next one year.
* **Modify**: Make modifications to the generated roster.
""")

uploaded_file = st.file_uploader("Upload current month's roster", type=['xlsx'])

if uploaded_file:
    current_roster = parse_roster(uploaded_file)
    st.write("### Current Month's Roster")
    st.dataframe(current_roster)

    generate_button = st.button("Generate One-Year Roster")
    if generate_button:
        one_year_roster = generate_yearly_roster(current_roster)
        st.write("### One-Year Duty Roster")
        st.dataframe(one_year_roster)

        # Option to download the generated roster
        @st.cache
        def convert_df(df):
            return df.to_csv(index=False).encode('utf-8')

        csv = convert_df(one_year_roster)

        st.download_button(
            label="Download roster as CSV",
            data=csv,
            file_name='one_year_roster.csv',
            mime='text/csv',
        )

# Save changes and regenerate
st.write("### Modify and Regenerate Roster")
edit_roster = st.text_area("Edit the roster here (CSV format):", height=300)

if st.button("Save Changes and Regenerate"):
    edited_df = pd.read_csv(pd.compat.StringIO(edit_roster))
    one_year_roster = generate_yearly_roster(edited_df)
    st.write("### Updated One-Year Duty Roster")
    st.dataframe(one_year_roster)