import plotly.io as pio
pio.renderers.default = 'browser'
import plotly.graph_objects as go
import numpy as np

z_base = np.array([
    [10, 10.625, 12.5, 15.625, 20],
    [5.625, 6.25, 8.125, 11.25, 15.625],
    [2.5, 3.125, 5., 8.125, 12.5],
    [0.625, 1.25, 3.125, 6.25, 10.625],
    [0, 0.625, 2.5, 5.625, 10]
])

lk_x=[0,0,1,1,0]
lk_y=[0,1,1,0,0]

data_line = go.Scatter(x=lk_x, y=lk_y, mode='lines', fill='toself', fillcolor='rgba(60, 53, 48, 0.5)', line_color='black', showlegend=False)

scales = np.linspace(0.2, 1.0, 10)
frames = []
for i, s in enumerate(scales):
    z_scaled = z_base * s
    frames.append(go.Frame(
        data=[
            go.Contour(
                z=z_scaled,
                colorbar=dict(
                    thickness=25,
                    thicknessmode='pixels',
                    len=0.6,
                    lenmode='fraction',
                    outlinewidth=0
                )
                ,zorder=2
            ),
            go.Scatter(
                x=[0, 1, 2, 3, 4],
                y=[0, 1, 1.5, 2, 3], 
                mode='lines',
                line=dict(color='black', width=2),
                showlegend=False,
                zorder=0
            ),
            data_line
        ],
        name=f"scale_{i}"
    ))


fig = go.Figure(
    data=frames[0].data,
    frames=frames
)

fig.update_layout(
    title="Contour",
    updatemenus=[{
        "buttons": [
            {"label": "play", "method": "animate", "args": [None, {"frame": {"duration": 500}, "fromcurrent": True}]},
            {"label": "pause", "method": "animate", "args": [[None], {"mode": "immediate", "frame": {"duration": 0}}]}
        ],
        "type": "buttons",
        "showactive": True,
        "x": 0.1,
        "y": -0.1
    }],
    sliders=[{
        "steps": [{"args": [[f.name]], "label": f"Step {i+1}", "method": "animate"} for i, f in enumerate(frames)],
        "x": 0.1,
        "len": 0.9
    }],
    width=500,
    height=500
)

fig.show()