import numpy as np
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import streamlit as st

def calculate_wave(wavelength, amplitude, frequency, num_periods, time_step):
    k = 2*np.pi/wavelength
    omega = 2*np.pi*frequency
    period = 1/frequency
    time_array = np.arange(0, num_periods*period, time_step)
    wave = amplitude * np.sin(k * np.arange(0, 1, wavelength/1000)[:, None] - omega * time_array)
    return time_array, wave

def plot_3d_wave(wavelength, amplitude, frequency):
    fig = make_subplots(rows=1, cols=2,
                        specs=[[{'type': 'surface'}, {'type': 'surface'}]],
                        subplot_titles=('Electric field', 'Magnetic field'),
                        )
    time_step = 1/(100*frequency)
    num_periods = 5
    time_array, wave = calculate_wave(wavelength, amplitude, frequency, num_periods, time_step)
    E_x = wave*np.cos(2*np.pi*time_array*frequency)
    E_y = np.zeros_like(E_x)
    E_z = np.zeros_like(E_x)
    B_x = np.zeros_like(E_x)
    B_y = -wave*np.sin(2*np.pi*time_array*frequency)
    B_z = np.zeros_like(E_x)
    
    fig.add_trace(go.Surface(x=wave, y=E_x, z=E_y, colorscale='Blues'), row=1, col=1)
    fig.add_trace(go.Surface(x=wave, y=B_x, z=B_y, colorscale='Reds'), row=1, col=2)
    fig.update_layout(scene_aspectratio=dict(x=1, y=1, z=1), width=800, height=400,
                      scene=dict(xaxis_title='Wave', yaxis_title='E_x', zaxis_title='E_y'),
                      scene2=dict(xaxis_title='Wave', yaxis_title='B_x', zaxis_title='B_y'),
                     )
    
    fig.update_xaxes(range=[0, wavelength], row=1, col=1)
    fig.update_xaxes(range=[0, wavelength], row=1, col=2)
    fig.update_yaxes(range=[-amplitude, amplitude], row=1, col=1)
    fig.update_yaxes(range=[-amplitude, amplitude], row=1, col=2)
    fig.update_zaxes(range=[-amplitude, amplitude], row=1, col=1)
    fig.update_zaxes(range=[-amplitude, amplitude], row=1, col=2)
    
    return fig

wavelength = st.slider('Wavelength', 0.1, 10.0, 1.0)
amplitude = st.slider('Amplitude', 0.1, 1.0, 0.5)
frequency = st.slider('Frequency', 0.1, 10.0, 1.0)

fig = plot_3d_wave(wavelength, amplitude, frequency)

st.plotly_chart(fig)