import json
from datetime import date, datetime
from urllib.request import urlopen

from matplotlib import dates as mdates
from matplotlib import pyplot as plt
from matplotlib.ticker import FuncFormatter

# Download the raw statistics from GitHub.
base_url = "https://raw.githubusercontent.com/cobyqa/stats/main/archives/"
pypi = json.loads(urlopen(base_url + "pypi.json").read())
conda = json.loads(urlopen(base_url + "conda.json").read())

# Extract the download statistics by excluding the mirror downloads.
pypi = {d["date"]: d["downloads"] for d in pypi if d["category"] == "without_mirrors"}
conda = {d["date"]: d["downloads"] for d in conda}

# Get the download dates.
pypi_download_dates = [datetime.strptime(d, "%Y-%m-%d").date() for d in pypi]
conda_download_dates = [datetime.strptime(d, "%Y-%m-%d").date() for d in conda]
all_download_dates = sorted(set(pypi_download_dates + conda_download_dates))

# Compute the cumulative download statistics.
pypi_download_count = []
conda_download_count = []
for download_date in all_download_dates:
    if download_date in pypi_download_dates:
        pypi_download_count.append(pypi[download_date.strftime("%Y-%m-%d")])
    else:
        pypi_download_count.append(0)
    if download_date in conda_download_dates:
        conda_download_count.append(conda[download_date.strftime("%Y-%m-%d")])
    else:
        conda_download_count.append(0)
pypi_download_cumulative = [sum(pypi_download_count[:i]) for i in range(1, len(pypi_download_count) + 1)]
conda_download_cumulative = [sum(conda_download_count[:i]) for i in range(1, len(conda_download_count) + 1)]

# Plot the cumulative downloads.
fig, ax = plt.subplots()
ax.xaxis.set_minor_locator(mdates.MonthLocator())
ax.xaxis.set_major_locator(mdates.YearLocator())
ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
ax.yaxis.set_major_formatter(FuncFormatter(lambda y, _: f"{int(y):,}"))
ax.plot(all_download_dates, pypi_download_cumulative, label="PyPI")
ax.plot(all_download_dates, conda_download_cumulative, label="conda-forge")
ax.set_xlim(date(2023, 1, 9), all_download_dates[-1])
ax.set_ylim(0, pypi_download_cumulative[-1] + conda_download_cumulative[-1])
ax.legend(loc="upper left")
ax.set_title("Cumulative downloads of COBYQA")