import streamlit as st
import numpy as np
from scipy.integrate import odeint
import plotly.graph_objects as go
import time
import pandas as pd

# Define the Lorenz system
def lorenz_system(current_state, t, sigma, rho, beta):
    x, y, z = current_state
    dx_dt = sigma * (y - x)
    dy_dt = x * (rho - z) - y
    dz_dt = x * y - beta * z
    return [dx_dt, dy_dt, dz_dt]

# Set up the sidebar
st.sidebar.title('Lorenz System Parameters')
sigma = st.sidebar.slider('Sigma', 0.0, 50.0, 10.0)
rho = st.sidebar.slider('Rho', 0.0, 50.0, 28.0)
beta = st.sidebar.slider('Beta', 0.0, 50.0, 2.67)

x0_1 = st.sidebar.number_input('Initial x for 1st system', value=1.0)
y0_1 = st.sidebar.number_input('Initial y for 1st system', value=1.0)
z0_1 = st.sidebar.number_input('Initial z for 1st system', value=1.0)

x0_2 = st.sidebar.number_input('Initial x for 2nd system', value=1.0)
y0_2 = st.sidebar.number_input('Initial y for 2nd system', value=1.0)
z0_2 = st.sidebar.number_input('Initial z for 2nd system', value=1.0)

# Define the time points for the simulation
t = np.linspace(0, 4, 1000)

# Solve the Lorenz system
solution_1 = odeint(lorenz_system, (x0_1, y0_1, z0_1), t, args=(sigma, rho, beta))
solution_2 = odeint(lorenz_system, (x0_2, y0_2, z0_2), t, args=(sigma, rho, beta))

# Create a 3D plot of the solutions
plot_slot = st.empty()

stop = st.checkbox('Stop')

i = 0
while i < len(t) and not stop:
    fig = go.Figure(data=[
        go.Scatter3d(x=solution_1[:i, 0], y=solution_1[:i, 1], z=solution_1[:i, 2], mode='lines', line=dict(color='red'), name='System 1'),
        go.Scatter3d(x=solution_2[:i, 0], y=solution_2[:i, 1], z=solution_2[:i, 2], mode='lines', line=dict(color='blue'), name='System 2')
    ])

    fig.update_layout(scene=dict(xaxis=dict(range=[min(min(solution_1[:, 0]), min(solution_2[:, 0])), max(max(solution_1[:, 0]), max(solution_2[:, 0]))]),
                                 yaxis=dict(range=[min(min(solution_1[:, 1]), min(solution_2[:, 1])), max(max(solution_1[:, 1]), max(solution_2[:, 1]))]),
                                 zaxis=dict(range=[min(min(solution_1[:, 2]), min(solution_2[:, 2])), max(max(solution_1[:, 2]), max(solution_2[:, 2]))]),
                                 camera=dict(eye=dict(x=2*np.cos(i/10), y=2*np.sin(i/10), z=0.1))),
                     width=800, height=600)

    plot_slot.plotly_chart(fig)

    i += 1
    time.sleep(0.2)

if st.button('Export Data'):
    df_1 = pd.DataFrame(solution_1, columns=['x_1', 'y_1', 'z_1'])
    df_2 = pd.DataFrame(solution_2, columns=['x_2', 'y_2', 'z_2'])
    df_1.to_csv('lorenz_data_1.csv')
    df_2.to_csv('lorenz_data_2.csv')
    st.write('Data exported successfully!')