import matplotlib.pyplot as plt
from pyscript import display, fetch, window, storage, document
from pyodide.ffi.wrappers import add_event_listener
import numpy as np
import time
import js

ONE_BILLION = 10**9
ONE_TRILLION = 10**12
ONE_QUADRILLION = 10**15

cache = await storage("cache")

async def set_musk_net_worth(event):
    input = document.getElementById("input")
    cache["musk_net_worth"] = input.value
    js.window.location.reload()
    
set = js.document.getElementById("set")
add_event_listener(set, "click", set_musk_net_worth)

async def reset_musk_net_worth(event):
    cache["musk_net_worth"] = None
    js.window.location.reload()
    
reset = js.document.getElementById("reset")
add_event_listener(reset, "click", reset_musk_net_worth)

try:
    int(cache["musk_net_worth"])
    musk_net_worth = float(cache["musk_net_worth"])
    if not (0 <= musk_net_worth <= ONE_QUADRILLION):
        raise RuntimeError("Outside limits.")
    input_msg = "You manually set the net worth previously, so it may not be right. You can reset if you want."
except:
    musk_net_worth = ONE_TRILLION * 1.1
    input_msg = "Musk's real net worth, as of Jun 13, 2026."

try:
    cache["countries_net_worth"]
    if (cache["countries_net_worth"][0] + 10800) <= time.time():
        window.console.log("Cache is outdated.")
        raise RuntimeError("Cache is outdated.")
    window.console.log("Retrieved data from cache.")
except:
    response = await fetch("https://statisticsoftheworld.com/api/v2/indicator/IMF.NGDPD")
    response_json = await response.json()
    cache["countries_net_worth"] = (time.time(), response_json)
    window.console.log("Fetched data and cached it.")
countries_net_worth = cache["countries_net_worth"][1]["data"]

x_list, y_list = ["MUSK"], [musk_net_worth / ONE_BILLION]

end_count = 0
for item in countries_net_worth:
    if item["value"] < musk_net_worth:
        x_list.append(item["countryId"])
        y_list.append(item["value"] / ONE_BILLION)
        end_count += 1
    if end_count == 10:
        break

if len(x_list) < 10:
    x_list, y_list = ["MUSK"], [musk_net_worth / ONE_BILLION]
    for item in countries_net_worth[:10]:
        x_list.append(item["countryId"])
        y_list.append(item["value"] / ONE_BILLION)
      
fig, ax = plt.subplots()
ax.bar(x_list, y_list)

ax.set_xlabel("Country Code / Musk")
ax.set_ylabel("Annual GDP / Net Worth (Billions, USD)")
ax.set_title(f"Countries' Annual Output vs. Musk's Net Worth")

display(fig, target="mpl")

input_thing = document.getElementById("input")
input_thing.value = musk_net_worth
input_thing.removeAttribute("disabled")

input_type = document.getElementById("input_type")
input_type.textContent = input_msg

set = document.getElementById("set")
set.removeAttribute("disabled")

reset = document.getElementById("reset")
reset.removeAttribute("disabled")

loading = document.getElementById("loading")
loading.textContent = ""