import streamlit as st
import plotly.graph_objects as go
import random

# Initialize session state
if 'initialized' not in st.session_state:
    st.session_state.initialized = False
    st.session_state.condition_counts = {}
    st.session_state.total_treatments = 0
    st.session_state.successful_treatments = 0
    st.session_state.unsuccessful_treatments = 0
    st.session_state.patients_waiting = 0
    st.session_state.doctor_actions = {}
    st.session_state.penalties = {}
    st.session_state.feedback = {}

# Function to create gauge chart
def create_gauge(title, value, max_value, domain, gauge_color):
    fig = go.Figure()
    fig.add_trace(go.Indicator(
        mode="gauge+number",
        value=value,
        title={'text': title},
        domain=domain,
        gauge={'axis': {'range': [0, max_value]}, 'bar': {'color': gauge_color}}
    ))
    fig.update_layout(template='plotly_dark')
    return fig

# Generate random patient condition distribution
def distribute_patients():
    remaining_patients = 200
    conditions = ["Healthy", "Mild Illness", "Chronic Illness", "Emergency"]
    condition_counts = {condition: 0 for condition in conditions}
    for condition in conditions:
        count = random.randint(0, remaining_patients)
        condition_counts[condition] = count
        remaining_patients -= count
    condition_counts[conditions[-1]] += remaining_patients
    return condition_counts

# Generate random treatment metrics
def generate_treatment_metrics():
    total_treatments = random.randint(50, 150)
    successful_treatments = random.randint(0, total_treatments)
    unsuccessful_treatments = total_treatments - successful_treatments
    patients_waiting = 200 - total_treatments
    return total_treatments, successful_treatments, unsuccessful_treatments, patients_waiting

# Generate random doctor actions
def generate_doctor_actions(total_treatments):
    actions = ["Prescribe Medication", "Recommend Tests", "Consult Clinician", "Schedule Surgery"]
    action_counts = {action: 0 for action in actions}
    for _ in range(total_treatments):
        action = random.choice(actions)
        action_counts[action] += 1
    return action_counts

# Generate random penalties
def generate_penalties(total_treatments):
    penalties = ["Wrong Medication", "Missed Diagnosis", "Incorrect Test Recommendation", "Stress-Induced Mistake"]
    penalty_counts = {penalty: 0 for penalty in penalties}
    for _ in range(total_treatments // 6):
        penalty = random.choice(penalties)
        penalty_counts[penalty] += 1
    return penalty_counts

# Generate random patient feedback
def generate_feedback(total_treatments):
    feedback = ["Completely Recovered", "Further Treatment Needed", "Complication"]
    feedback_counts = {fb: 0 for fb in feedback}
    for _ in range(total_treatments):
        fb = random.choice(feedback)
        feedback_counts[fb] += 1
    return feedback_counts

# Streamlit setup
st.title("Interactive Healthcare Metrics")

# Button to run the animations
if st.button("Run Animations"):
    st.session_state.condition_counts = distribute_patients()
    st.session_state.total_treatments, st.session_state.successful_treatments, st.session_state.unsuccessful_treatments, st.session_state.patients_waiting = generate_treatment_metrics()
    st.session_state.doctor_actions = generate_doctor_actions(st.session_state.total_treatments)
    st.session_state.penalties = generate_penalties(st.session_state.total_treatments)
    st.session_state.feedback = generate_feedback(st.session_state.total_treatments)
    st.session_state.initialized = True

if st.session_state.initialized:
    # Slider to choose the animation
    animation_index = st.slider("Select Animation", min_value=1, max_value=5, value=1)

    # Animation 1: Patient Conditions
    if animation_index == 1:
        st.subheader("Patients by Condition")
        for condition, count in st.session_state.condition_counts.items():
            fig = create_gauge(f"Patients with {condition}", count, 200, {'x': [0, 1], 'y': [0, 1]}, 'blue')
            st.plotly_chart(fig, use_container_width=True)
            st.markdown("---")

    # Animation 2: Treatment Metrics
    elif animation_index == 2:
        st.subheader("Treatment Metrics")
        fig1 = create_gauge("Total Patients", 200, 200, {'x': [0, 0.5], 'y': [0, 1]}, 'green')
        fig2 = create_gauge("Successful Treatments", st.session_state.successful_treatments, st.session_state.total_treatments, {'x': [0.5, 1], 'y': [0, 0.5]}, 'orange')
        fig3 = create_gauge("Unsuccessful Treatments", st.session_state.unsuccessful_treatments, st.session_state.total_treatments, {'x': [0.5, 1], 'y': [0.5, 1]}, 'red')
        st.plotly_chart(fig1, use_container_width=True)
        st.plotly_chart(fig2, use_container_width=True)
        st.plotly_chart(fig3, use_container_width=True)
        st.markdown("---")

    # Animation 3: Doctor Actions
    elif animation_index == 3:
        st.subheader("Doctor Actions")
        for action, count in st.session_state.doctor_actions.items():
            fig = create_gauge(f"Doctor Action: {action}", count, st.session_state.total_treatments, {'x': [0, 1], 'y': [0, 1]}, 'purple')
            st.plotly_chart(fig, use_container_width=True)
            st.markdown("---")

    # Animation 4: Penalties
    elif animation_index == 4:
        st.subheader("Penalties")
        for penalty, count in st.session_state.penalties.items():
            fig = create_gauge(f"Penalty: {penalty}", count, st.session_state.total_treatments // 6, {'x': [0, 1], 'y': [0, 1]}, 'red')
            st.plotly_chart(fig, use_container_width=True)
            st.markdown("---")

    # Animation 5: Patient Feedback
    elif animation_index == 5:
        st.subheader("Patient Feedback")
        for fb, count in st.session_state.feedback.items():
            fig = create_gauge(f"Feedback: {fb}", count, st.session_state.total_treatments, {'x': [0, 1], 'y': [0, 1]}, 'blue')
            st.plotly_chart(fig, use_container_width=True)
            st.markdown("---")

    st.write("**End of Animations.**")