# -*- coding: utf-8 -*-
fig_pyr = go.Figure()

# Initial bars (first year)
fig_pyr.add_bar(y=[str(years[0])], x=[-mex_vals[0]],
                orientation="h", name="Mexico", marker_color="steelblue")
fig_pyr.add_bar(y=[str(years[0])], x=[ pol_vals[0]],
                orientation="h", name="Poland", marker_color="indianred")

frames = []
for i, yr in enumerate(years):
    frames.append(go.Frame(
        name=str(yr),
        data=[
            go.Bar(y=[str(yr)], x=[-mex_vals[i]]),
            go.Bar(y=[str(yr)], x=[ pol_vals[i]])
        ]
    ))
fig_pyr.frames = frames

fig_pyr.update_layout(
    title="Annual Head-to-Head Comparison of Tests per Thousand Inhabitants\n(Back-to-Back Pyramid)",
    xaxis=dict(
        range=[-max_abs, max_abs],
        tickvals=[-max_abs, -max_abs/2, 0, max_abs/2, max_abs],
        ticktext=[f"{abs(t):.0f}" for t in [-max_abs, -max_abs/2, 0, max_abs/2, max_abs]],
        title="Tests per Thousand (absolute value)"
    ),
    yaxis=dict(title="Year", autorange="reversed"),
    barmode="overlay",
    legend_title_text="Country",
    updatemenus=[{
        "type": "buttons",
        "showactive": False,
        "y": 1.13, "x": 1, "xanchor": "right",
        "buttons": [
            {"label": "▶", "method": "animate",
             "args": [None,
                      {"frame": {"duration": 600, "redraw": True},
                       "transition": {"duration": 300},
                       "fromcurrent": True}]},
            {"label": "⏸", "method": "animate",
             "args": [[None],
                      {"frame": {"duration": 0, "redraw": False},
                       "mode": "immediate"}]}
        ]
    }],
    sliders=[{
        "steps": [
            {"label": str(yr), "method": "animate",
             "args": [[str(yr)],
                      {"frame": {"duration": 0, "redraw": True},
                       "mode": "immediate"}]}
            for yr in years
        ],
        "currentvalue": {"prefix": "Year: "}
    }]
)

fig_pyr.show()

import pandas as pd
df=pd.read_csv("https://catalog.ourworldindata.org/garden/covid/latest/cases_deaths/cases_deaths.csv")

import numpy as np

df=df[df['country']=='Mexico']

df=df.reset_index(drop=True)

df.columns

df['weekly_avg'] = df['new_cases'].rolling(window=7).mean()

df['weekly_avg'] = df['weekly_avg'].fillna(0)

import plotly.graph_objects as go


df['date'] = pd.to_datetime(df['date'])
df = df.sort_values('date')
df['weekly_avg'] = df['new_cases'].rolling(window=7).mean()