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


# Function for disease task
def run_disease_task():
    # Number of samples per class
    N_per_class = 500

    # Number of classes
    num_classes = 5

    # Total number of samples
    N = N_per_class * num_classes

    # Number of features
    D = 2  # For visualization purposes

    # Initialize feature matrix X and label vector y
    X = np.zeros((N, D))
    y = np.zeros(N, dtype=int)

    # Generate a multi-class spiral dataset
    def generate_multi_class_spiral(points, classes):
        X = np.zeros((points * classes, 2))
        y = np.zeros(points * classes, dtype=int)
        for class_number in range(classes):
            ix = range(points * class_number, points * (class_number + 1))
            r = np.linspace(0.0, 1, points)
            t = np.linspace(class_number * 4, (class_number + 1) * 4, points) + np.random.randn(points) * 0.2
            X[ix] = np.c_[r * np.sin(t), r * np.cos(t)]
            y[ix] = class_number
        return X, y

    X, y = generate_multi_class_spiral(N_per_class, num_classes)

    # Shuffle the dataset
    X, y = shuffle(X, y, random_state=42)

    # Normalize the features
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # Convert data to torch tensors
    X_train_tensor = torch.from_numpy(X_scaled).float()
    y_train_tensor = torch.from_numpy(y).long()

    # Split data into training and test sets
    X_train_tensor, X_test_tensor, y_train_tensor, y_test_tensor = train_test_split(
        X_train_tensor, y_train_tensor, test_size=0.2, random_state=42
    )

    # Logistic Regression Model
    linear_model = LogisticRegression(max_iter=200)
    linear_model.fit(X_scaled[: int(0.8 * N)], y[: int(0.8 * N)])

    # Linear model accuracy
    linear_accuracy = linear_model.score(X_scaled[int(0.8 * N) :], y[int(0.8 * N) :])

    # Neural Network Model
    class NeuralNet(nn.Module):
        def __init__(self, input_dim, hidden_dims, output_dim):
            super(NeuralNet, self).__init__()
            layers = []
            in_dim = input_dim
            for h_dim in hidden_dims:
                layers.append(nn.Linear(in_dim, h_dim))
                layers.append(nn.ReLU())
                layers.append(nn.BatchNorm1d(h_dim))
                layers.append(nn.Dropout(0.3))
                in_dim = h_dim
            layers.append(nn.Linear(in_dim, output_dim))
            self.model = nn.Sequential(*layers)

        def forward(self, x):
            return self.model(x)

    # Initialize Neural Network Model
    hidden_dims = [128, 64, 32]
    neural_model = NeuralNet(D, hidden_dims, num_classes)

    # Loss and optimizer for Neural Network Model
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(neural_model.parameters(), lr=0.001, weight_decay=1e-4)

    # Training the Neural Network Model
    num_epochs = 200
    for epoch in range(num_epochs):
        neural_model.train()
        outputs = neural_model(X_train_tensor)
        loss = criterion(outputs, y_train_tensor)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (epoch + 1) % 20 == 0:
            st.write(f'Neural Network - Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

    # Evaluate Neural Network Model
    neural_model.eval()
    with torch.no_grad():
        outputs = neural_model(X_test_tensor)
        _, predicted = torch.max(outputs.data, 1)
        neural_accuracy = (predicted == y_test_tensor).sum().item() / y_test_tensor.size(0)
        st.write(f'Neural Network Model Accuracy: {neural_accuracy * 100:.2f}%')

    # Summary of Accuracies
    st.write("\nSummary of Accuracies:")
    st.write(f'Linear Model Accuracy: {linear_accuracy * 100:.2f}%')
    st.write(f'Neural Network Model Accuracy: {neural_accuracy * 100:.2f}%')


# 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"):
        run_disease_task()