import gradio as gr
import pandas as pd
import random
from faker import Faker
from pyspark.sql import SparkSession
from pyspark.sql.functions import avg, count, round as spark_round

fake = Faker()

spark = SparkSession.builder \
    .appName("SyntheticPatientAdmissions") \
    .master("local[*]") \
    .getOrCreate()


diagnoses = [
    "Pneumonia", "Diabetes", "Hypertension", "Asthma",
    "Heart Failure", "COVID-19", "Kidney Disease",
    "Fracture", "Migraine", "Sepsis"
]

departments = [
    "Emergency", "Cardiology", "Pulmonology",
    "Orthopedics", "Neurology", "ICU",
    "General Medicine"
]

insurance_types = [
    "Medicare", "Medicaid", "Private", "Self-Pay", "VA"
]


def generate_patient_data(num_records):
    records = []

    for i in range(int(num_records)):
        age = random.randint(1, 95)
        length_of_stay = random.randint(1, 21)
        admission_cost = round(random.uniform(1200, 45000), 2)

        records.append({
            "patient_id": f"P{i+1:05d}",
            "name": fake.name(),
            "age": age,
            "gender": random.choice(["Male", "Female"]),
            "admission_date": fake.date_between(start_date="-1y", end_date="today"),
            "diagnosis": random.choice(diagnoses),
            "department": random.choice(departments),
            "insurance_type": random.choice(insurance_types),
            "length_of_stay": length_of_stay,
            "admission_cost": admission_cost
        })

    return pd.DataFrame(records)


def analyze_patient_data(num_records):
    df = generate_patient_data(num_records)

    spark_df = spark.createDataFrame(df)

    avg_stay = spark_df.select(
        spark_round(avg("length_of_stay"), 2).alias("average_length_of_stay")
    ).toPandas()

    top_diagnoses = spark_df.groupBy("diagnosis") \
        .agg(count("*").alias("count")) \
        .orderBy("count", ascending=False) \
        .toPandas()

    admissions_by_department = spark_df.groupBy("department") \
        .agg(count("*").alias("admissions")) \
        .orderBy("admissions", ascending=False) \
        .toPandas()

    avg_cost_by_insurance = spark_df.groupBy("insurance_type") \
        .agg(spark_round(avg("admission_cost"), 2).alias("average_cost")) \
        .orderBy("average_cost", ascending=False) \
        .toPandas()

    age_summary = spark_df.select(
        spark_round(avg("age"), 2).alias("average_age")
    ).toPandas()

    return (
        df,
        avg_stay,
        top_diagnoses,
        admissions_by_department,
        avg_cost_by_insurance,
        age_summary
    )


with gr.Blocks(title="Synthetic Patient Admissions with PySpark") as demo:
    gr.Markdown(
        """
        # Synthetic Patient Admission Records
        
        This beginner-friendly app uses **Faker** to generate synthetic healthcare admission records 
        and **PySpark** to analyze the data.
        """
    )

    num_records = gr.Slider(
        minimum=10,
        maximum=1000,
        value=100,
        step=10,
        label="Number of Patient Records"
    )

    generate_button = gr.Button("Generate and Analyze Data")

    patient_table = gr.Dataframe(label="Synthetic Patient Admission Records")
    avg_stay_table = gr.Dataframe(label="Average Length of Stay")
    diagnosis_table = gr.Dataframe(label="Most Common Diagnoses")
    department_table = gr.Dataframe(label="Admissions by Department")
    insurance_table = gr.Dataframe(label="Average Cost by Insurance Type")
    age_table = gr.Dataframe(label="Average Patient Age")

    generate_button.click(
        fn=analyze_patient_data,
        inputs=num_records,
        outputs=[
            patient_table,
            avg_stay_table,
            diagnosis_table,
            department_table,
            insurance_table,
            age_table
        ]
    )

demo.launch()