import streamlit as st
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.utils import shuffle
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

from NdR_disease import run_disease_train, get_user_input_and_predict


# Function for male superhero task
def run_male_superhero_task():
    st.write("Training Male Superhero model...")
    # Male superhero training logic goes here
    # Add dummy print statements as a placeholder
    st.write("Male superhero model - Step 1: Data prepared.")
    st.write("Male superhero model - Step 2: Model trained.")
    st.write("Male superhero model - Step 3: Results evaluated.")


# Function for female superhero task
def run_female_superhero_task():
    st.write("Training Female Superhero model...")
    # Female superhero training logic goes here
    # Add dummy print statements as a placeholder
    st.write("Female superhero model - Step 1: Data prepared.")
    st.write("Female superhero model - Step 2: Model trained.")
    st.write("Female superhero model - Step 3: Results evaluated.")


# Streamlit UI
st.title("AI Training Demo")

# Task selection buttons
task = st.selectbox("Choose a task:", ("Superhero", "Disease"))

if task == "Superhero":
    # Sub-options for Male and Female Superhero
    gender = st.selectbox("Choose the gender:", ("Male", "Female"))

    if gender == "Male":
        if st.button("Run Male Superhero Task"):
            run_male_superhero_task()

    elif gender == "Female":
        if st.button("Run Female Superhero Task"):
            run_female_superhero_task()

elif task == "Disease":
    if st.button("Run Disease Task"):
        linear_model, neural_model, scaler, conditions, num_classes = run_disease_train()
        get_user_input_and_predict(linear_model, neural_model, scaler, conditions, num_classes)