# -*- coding: utf-8 -*-
plt.figure(figsize=(14, 7))
for country in ["Mexico", "Poland"]:
    subset = df[df["country"] == country]
    plt.plot(
        subset["date"],
        subset["total_tests_per_thousand"],
        linewidth=2,
        label=country
    )


plt.title("Daily Evolution of COVID-19 Tests per Thousand: Mexico vs Poland", fontsize=14)
plt.xlabel("Date", fontsize=12)
plt.ylabel("Tests per Thousand", fontsize=12)
plt.xticks(rotation=45)
plt.legend(title="Country")
plt.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()

# Format x-axis with months
plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
plt.gca().xaxis.set_major_locator(mdates.MonthLocator(interval=3))

# Set limits
min_date = df["date"].min()
max_date = df["date"].max()
plt.xlim([min_date, max_date])

plt.show()

import pandas as pd
import matplotlib.pyplot as plt


df = pd.read_csv("testing.csv", parse_dates=["date"])
df = (
    df[df["country"].isin(["Mexico", "Poland"])]
      .dropna(subset=["total_tests_per_thousand"])
      .copy()
)

df["year"] = df["date"].dt.year
annual = (
    df.sort_values("date")
      .groupby(["country", "year"], as_index=False)
      .last()[["country", "year", "total_tests_per_thousand"]]
)

pivot_annual = (
    annual.pivot(index="year", columns="country", values="total_tests_per_thousand")
          .fillna(0)
)

# Data
years       = pivot_annual.index.astype(str)
mex_values  = pivot_annual["Mexico"]
pol_values  = pivot_annual["Poland"]

fig, ax = plt.subplots(figsize=(12, 8))

# Horizontal bars
bars_mex = ax.barh(years, -mex_values, label="Mexico", color='tab:red')
bars_pol = ax.barh(years,  pol_values, label="Poland", color='tab:blue')


max_val = max(mex_values.max(), pol_values.max())
x_max = int((max_val + 5) // 10 * 10 + 10)

# grid
tick_step = 10
label_step = 80
all_ticks = list(range(-x_max, x_max + 1, tick_step))

ax.set_xlim(-x_max, x_max)
ax.set_xticks(all_ticks)
ax.set_xticklabels([
    str(abs(x)) if x % label_step == 0 or x == 0 else "" for x in all_ticks
], rotation=45, ha='right', fontsize=9)

for bar in bars_mex:
    width = bar.get_width()
    ax.text(width - 1, bar.get_y() + bar.get_height() / 2,
            f'{abs(width):.1f}', va='center', ha='right', fontsize=9)

for bar in bars_pol:
    width = bar.get_width()
    ax.text(width + 1, bar.get_y() + bar.get_height() / 2,
            f'{width:.1f}', va='center', ha='left', fontsize=9)

ax.set_xlabel("Total Tests per Thousand (absolute value)", fontsize=12)
ax.set_title("Annual Head-to-Head Comparison of COVID-19 Testing: Mexico vs Poland", fontsize=14)
ax.axvline(0, color='gray', linestyle='--', linewidth=1)
ax.legend()
ax.grid(True, axis='x', which='both', linestyle='--', alpha=0.4)
plt.tight_layout()
plt.show()

"""### **ANIMATED**"""

import pandas as pd
import plotly.graph_objects as go

df = pd.read_csv("testing.csv", parse_dates=["date"])

df = (
    df[df["country"].isin(["Mexico", "Poland"])]
      .dropna(subset=["total_tests_per_thousand"])
      .copy()
)


df["period"] = df["date"].dt.to_period("M").astype(str)

monthly = (
    df.groupby(["country", "period"], as_index=False)
      .last()           # last record of each month
      .sort_values(["period", "country"])
)

periods  = monthly["period"].unique()
countries = ["Mexico", "Poland"]