import streamlit as st
import numpy as np
import plotly.graph_objects as go
from streamlit_plotly_events import plotly_events

st.set_page_config(page_title="Dot Product Visualization", layout="wide")

# Initialize session state
if "v1" not in st.session_state:
    st.session_state.v1 = np.array([3, 2])  # Initial position of vector 1
if "v2" not in st.session_state:
    st.session_state.v2 = np.array([-2, 3])  # Initial position of vector 2
if "selected_vector" not in st.session_state:
    st.session_state.selected_vector = None  # To track which vector is being moved

# Function to compute dot product
def compute_dot_product(v1, v2):
    return np.dot(v1, v2)

# Create a plot
fig = go.Figure()

# Add vector 1
fig.add_trace(go.Scatter(
    x=[0, st.session_state.v1[0]], 
    y=[0, st.session_state.v1[1]], 
    mode="lines+markers", 
    marker=dict(size=12, color="blue"),
    line=dict(width=4, color="blue"),
    name="Vector 1"
))

# Add vector 2
fig.add_trace(go.Scatter(
    x=[0, st.session_state.v2[0]], 
    y=[0, st.session_state.v2[1]], 
    mode="lines+markers", 
    marker=dict(size=12, color="red"),
    line=dict(width=4, color="red"),
    name="Vector 2"
))

# Set axis properties
fig.update_layout(
    xaxis=dict(range=[-5, 5], title="X-axis"),
    yaxis=dict(range=[-5, 5], title="Y-axis"),
    title="Interactive Dot Product Visualization",
    dragmode="pan",
    showlegend=True
)

# Display the interactive plot
st.subheader(f"Dot Product: `{compute_dot_product(st.session_state.v1, st.session_state.v2):.2f}`")
selected_points = plotly_events(fig, click_event=True)

# Check if a point is clicked
if selected_points:
    clicked_x, clicked_y = selected_points[0]['x'], selected_points[0]['y']

    # Determine which vector is clicked
    if np.allclose([clicked_x, clicked_y], st.session_state.v1, atol=0.5):
        st.session_state.selected_vector = "v1"
    elif np.allclose([clicked_x, clicked_y], st.session_state.v2, atol=0.5):
        st.session_state.selected_vector = "v2"

# Move the selected vector using sliders
if st.session_state.selected_vector:
    st.sidebar.subheader(f"Move {st.session_state.selected_vector}")
    new_x = st.sidebar.slider("X component", -5.0, 5.0, float(st.session_state[st.session_state.selected_vector][0]), step=0.1)
    new_y = st.sidebar.slider("Y component", -5.0, 5.0, float(st.session_state[st.session_state.selected_vector][1]), step=0.1)

    # Update vector position
    st.session_state[st.session_state.selected_vector] = np.array([new_x, new_y])