import streamlit as st
import numpy as np
import torch
import torch.nn as nn
import random

# Define the function to optimize
def f(x):
    return x**2

# Define the neural network
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(1, 10)
        self.fc2 = nn.Linear(10, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Define the genetic algorithm
class GeneticAlgorithm:
    def __init__(self, population_size):
        self.population_size = population_size
        self.population = [random.uniform(-10, 10) for _ in range(population_size)]

    def selection(self):
        self.population = sorted(self.population, key=lambda x: f(x), reverse=True)[:self.population_size//2]

    def crossover(self):
        offspring = []
        for _ in range(self.population_size//2):
            parent1, parent2 = random.sample(self.population, 2)
            child = (parent1 + parent2) / 2
            offspring.append(child)
        self.population += offspring

    def mutation(self):
        for i in range(self.population_size):
            if random.random() < 0.1:
                self.population[i] += random.uniform(-1, 1)

# Streamlit app
st.title("Neural Network vs Genetic Algorithm")

# Parameters
st.sidebar.header("Parameters")
population_size = st.sidebar.slider("Population size", 10, 100, 50)
nn_iterations = st.sidebar.slider("Neural network iterations", 10, 1000, 100)
ga_iterations = st.sidebar.slider("Genetic algorithm iterations", 10, 100, 50)

# Run the experiment
if st.button("Run experiment"):
    # Initialize the neural network and genetic algorithm
    net = Net()
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
    ga = GeneticAlgorithm(population_size)

    # Train the neural network
    nn_losses = []
    for i in range(nn_iterations):
        x = torch.randn(1, 1)
        y = f(x)
        y_pred = net(x)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        nn_losses.append(loss.item())

    # Run the genetic algorithm
    ga_values = []
    for i in range(ga_iterations):
        ga.selection()
        ga.crossover()
        ga.mutation()
        ga_values.append(max(map(f, ga.population)))

    # Plot the results
    st.line_chart({"Neural Network": nn_losses, "Genetic Algorithm": ga_values})

    # Print the final values
    st.write("Final neural network value:", nn_losses[-1])
    st.write("Final genetic algorithm value:", ga_values[-1])