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_disease
from NdR_female_superheros import run_female_superhero_train,  get_user_input_and_predict_female_superhero
from NdR_male_superheros import run_male_superhero_train,  get_user_input_and_predict_male_superhero


# Interfaccia utente Streamlit
st.title("Demo: Intelligenza Artificiale")

# Pulsanti di selezione del attività
task = st.selectbox("Scegli una attività:", ("Supereroe", "Malattie"))

if task == "Supereroe":
    # Sotto-opzioni per Supereroe Maschile e Femminile
    gender = st.selectbox("Scegli il genere:", ("Maschile", "Femminile"))

    if gender == "Maschile":
        if 'malesuper_models' not in st.session_state:
            if st.button("Avvia attività Supereroe Maschile"):
                # Carica i modelli e salvali nello stato di sessione
                linear_model, neural_model, scaler, superheros, num_classes = run_male_superhero_train()
                st.session_state.malesuper_models = {
                    'linear_model': linear_model,
                    'neural_model': neural_model,
                    'scaler': scaler,
                    'superheros': superheros,
                    'num_classes': num_classes
                }
        else:
            st.write("I modelli di Supereroe Maschile sono già caricati.")

        # Se i modelli sono caricati, ottieni input dall'utente e fai una previsione
        if 'malesuper_models' in st.session_state:
            models = st.session_state.malesuper_models
            get_user_input_and_predict_male_superhero(
                models['linear_model'],
                models['neural_model'],
                models['scaler'],
                models['superheros'],
                models['num_classes']
            )
        else:
            st.write("Clicca 'Avvia attività Supereroe Maschile' per caricare i modelli.")

    elif gender == "Femminile":
        if 'femalesuper_models' not in st.session_state:
            if st.button("Avvia attività Supereroe Femminile"):
                # Carica i modelli e salvali nello stato di sessione
                linear_model, neural_model, scaler, superheros, num_classes = run_female_superhero_train()
                st.session_state.femalesuper_models = {
                    'linear_model': linear_model,
                    'neural_model': neural_model,
                    'scaler': scaler,
                    'superheros': superheros,
                    'num_classes': num_classes
                }
        else:
            st.write("I modelli di Supereroe Femminile sono già caricati.")

        # Se i modelli sono caricati, ottieni input dall'utente e fai una previsione
        if 'femalesuper_models' in st.session_state:
            models = st.session_state.femalesuper_models
            get_user_input_and_predict_female_superhero(
                models['linear_model'],
                models['neural_model'],
                models['scaler'],
                models['superheros'],
                models['num_classes']
            )
        else:
            st.write("Clicca 'Avvia attività Supereroe Femminile' per caricare i modelli.")

elif task == "Malattie":
    if 'disease_models' not in st.session_state:
        if st.button("Avvia attività Malattie"):
            # Carica i modelli e salvali nello stato di sessione
            linear_model, neural_model, scaler, conditions, num_classes = run_disease_train()
            st.session_state.disease_models = {
                'linear_model': linear_model,
                'neural_model': neural_model,
                'scaler': scaler,
                'conditions': conditions,
                'num_classes': num_classes
            }
    else:
        st.write("I modelli di malattie sono già caricati.")

    # Se i modelli sono caricati, ottieni input dall'utente e fai una previsione
    if 'disease_models' in st.session_state:
        models = st.session_state.disease_models
        get_user_input_and_predict_disease(
            models['linear_model'],
            models['neural_model'],
            models['scaler'],
            models['conditions'],
            models['num_classes']
        )
    else:
        st.write("Clicca 'Avvia attività Malattie' per caricare i modelli.")

if st.button('Ripristina'):
    # Cancella solo lo stato di invio del modulo e gli input dell'utente
    for key in list(st.session_state.keys()):
        del st.session_state[key]
    st.rerun()