import ipyreact
from traitlets import List, Any
import numpy as np

class TldrawSineWidget(ipyreact.ReactWidget):
    points = List(List(Any())).tag(sync=True)
    _esm = """
    import { TDShapeType, Tldraw } from "@tldraw/tldraw";
    import * as React from "react";

    export default function App({ points }) {

        const [app, setApp] = React.useState()

        const handleMount = React.useCallback((app: Tldraw) => {
            setApp(app)
        }, []);

        React.useEffect(() => {
            if (app) {
            app.createShapes({
                type: "draw",
                id: "draw1",
                color: 'red',
                points: points,
            });
            }
        }, [points, app])

        return (
            <div
            style={{
                position: "relative",
                width: "800px",
                height: "450px",
            }}
            >
            <Tldraw onMount={handleMount} onChange={e => console.log("hii")} />
            </div>
        );
}

    """

import solara
import numpy as np

float_value = solara.reactive(0)

@solara.component
def Page():
    dispersion = "high" # "high", "low", "const"
    def n(wj):
        if dispersion == "high":
            return 1 + wj * 0.1
        if dispersion == "low":
            return 1.1 - wj * 0.01
        if dispersion == "const":
            return 1

    c = 1
    num_of_waves = 121
    start_w = 1
    end_w = 7
    x = np.linspace(-2, 10, 1001)
    
    def g(x, t):
        u1 = 0
        for wj in np.linspace(start_w, end_w, num_of_waves):
            u1 += np.exp(1j * (wj * n(wj) / c * x - t * wj))
        return u1.real

    ooo = 38    
    points = np.column_stack((x*50+ooo, g(x, float_value.value)+150)).tolist()
    points = points[100:]
    points[0] = [ooo-40, 0]


    # solara.Markdown(" \psi(x, t) = \sum \limits_j C_j \cdot e^{\mathrm{i}( - k_j \cdot x + \omega_j \cdot t )}, \omega_j= \frac{k_j}{c(k_j)} $")
    # solara.Markdown(f"**t**: {float_value.value}")


    with solara.Row():
        solara.Button("Reset", on_click=lambda: float_value.set(0))
        solara.FloatSlider("Time t", value=float_value, min=0, max=4*np.pi, step=0.1)
    TldrawSineWidget.element(points=points)



Page()