import streamlit as st
import pandas as pd
from collections import defaultdict
import plotly.express as px

# Title and file upload
st.title("Enhanced Duty Roster Generator with Role-Specific Assignments")
uploaded_file = st.file_uploader("Upload your availability spreadsheet with roles", type=["csv", "xlsx"])

# Multi-shift configuration
shifts = st.multiselect("Select shifts for each day", ["Morning", "Afternoon", "Night"], ["Morning", "Afternoon"])
shift_requirements = {shift: st.number_input(f"Personnel required for {shift} shift", min_value=1, value=1) for shift in shifts}

if uploaded_file:
    # Load the file
    if uploaded_file.name.endswith('.csv'):
        df = pd.read_csv(uploaded_file)
    else:
        df = pd.read_excel(uploaded_file)

    # Display data and allow user to confirm roles
    st.write("Personnel Availability and Roles Data:")
    st.dataframe(df)

    # Generate duty roster with roles
    def generate_duty_roster_with_roles(data, shifts, shift_reqs):
        roster = defaultdict(lambda: defaultdict(list))

        for day in range(1, 31):  # Example with 30 days
            for shift in shifts:
                role_counts = defaultdict(int)
                assigned_today = 0

                for _, row in data.iterrows():
                    name = row['Name']
                    role = row['Role']
                    unavailable_days = set(map(int, row['Unavailable Days'].split(',')))

                    if day not in unavailable_days and role_counts[role] < shift_reqs[shift]:
                        roster[day][shift].append((name, role))
                        role_counts[role] += 1
                        assigned_today += 1

                    if assigned_today == shift_reqs[shift]:
                        break

                # If unable to meet required personnel, mark as unassigned
                if assigned_today < shift_reqs[shift]:
                    roster[day][shift].append(("Unassigned", "N/A"))

        return roster

    # Generate and display the duty roster
    duty_roster = generate_duty_roster_with_roles(df, shifts, shift_requirements)

    # Display the roster with roles and shifts
    st.write("Generated Duty Roster with Roles:")
    for day, shifts in duty_roster.items():
        st.write(f"Day {day}:")
        for shift, personnel in shifts.items():
            st.write(f"{shift} Shift: {', '.join([f'{p[0]} ({p[1]})' for p in personnel])}")

    # Interactive Calendar Visualization
    flattened_data = []
    for day, shifts in duty_roster.items():
        for shift, personnel in shifts.items():
            for name, role in personnel:
                flattened_data.append({"Day": day, "Shift": shift, "Personnel": name, "Role": role})

    df_roster = pd.DataFrame(flattened_data)
    fig = px.timeline(df_roster, x_start="Day", x_end="Day", y="Personnel", color="Shift", title="Duty Roster Calendar View")
    st.plotly_chart(fig)

    # Downloadable individual schedules
    st.write("Download Individual Schedules:")
    for person in df['Name'].unique():
        person_schedule = df_roster[df_roster['Personnel'] == person]
        person_csv = person_schedule.to_csv(index=False).encode('utf-8')
        st.download_button(f"Download Schedule for {person}", data=person_csv, file_name=f"{person}_schedule.csv", mime="text/csv")

    # Scheduling Insights and Analytics
    st.write("Scheduling Insights:")
    personnel_counts = df_roster.groupby("Personnel").size()
    st.bar_chart(personnel_counts)
    st.write(f"Total shifts assigned per person: {personnel_counts.to_dict()}")

    st.write("Days with unassigned shifts:")
    unassigned_days = df_roster[df_roster['Personnel'] == "Unassigned"]
    st.write(unassigned_days[['Day', 'Shift']].drop_duplicates())