import altair as alt
import numpy as np
import pandas as pd
import streamlit as st

st.set_page_config(layout="wide")

st.title("Fast Spiral Visualization")

num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
num_turns = st.slider("Number of turns in spiral", 1, 300, 31)


@st.cache_data(show_spinner=False)
def generate_spiral(num_points, num_turns):
    indices = np.linspace(0, 1, num_points, dtype=np.float32)
    theta = 2 * np.pi * num_turns * indices

    x = indices * np.cos(theta)
    y = indices * np.sin(theta)

    return pd.DataFrame({
        "x": x,
        "y": y,
        "idx": indices,
    })


df = generate_spiral(num_points, num_turns)

chart = (
    alt.Chart(df, height=700, width=700)
    .mark_circle(size=25)   # fixed size = much faster
    .encode(
        x=alt.X("x:Q", axis=None),
        y=alt.Y("y:Q", axis=None),
        color=alt.Color(
            "idx:Q",
            legend=None,
            scale=alt.Scale(scheme="turbo")  # faster & nicer
        ),
    )
)

st.altair_chart(chart, use_container_width=True)