Update price shock page (not working)

requirements.txt

@@ -1,12 +1,14 @@
 aiodns==3.0.0
-aiohttp==3.10.2
+aiohttp==3.8.3
 aiosignal==1.3.1
 altair==5.3.0
 anchorpy==0.20.1
 anchorpy-core==0.2.0
 annotated-types==0.7.0
 anyio==3.6.2
 apischema==0.17.5
+appnope==0.1.4
+asttokens==2.4.1
 async-timeout==4.0.2
 attrs==22.1.0
 backoff==2.2.1
@@ -19,13 +21,17 @@ certifi==2022.12.7
 cffi==1.15.1
 charset-normalizer==2.1.1
 click==8.1.7
+comm==0.2.2
 construct==2.10.68
 construct-typing==0.5.3
+debugpy==1.8.6
+decorator==5.1.1
 Deprecated==1.2.14
 dnspython==2.2.1
 driftpy==0.7.89
 Events==0.5
 exceptiongroup==1.0.4
+executing==2.1.0
 fastapi==0.115.0
 fastapi-restful==0.6.0
 fastapi-utils==0.7.0
@@ -41,16 +47,22 @@ httpx==0.23.1
 humanize==4.10.0
 idna==3.4
 iniconfig==1.1.1
+ipykernel==6.29.5
+ipython==8.28.0
+jedi==0.19.1
 Jinja2==3.1.4
 jito_searcher_client==0.1.4
 jsonalias==0.1.1
 jsonrpcclient==4.0.3
 jsonrpcserver==5.0.9
 jsonschema==4.17.3
+jupyter_client==8.6.3
+jupyter_core==5.7.2
 loguru==0.6.0
 Markdown==3.6
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
+matplotlib-inline==0.1.7
 mccabe==0.7.0
 mdurl==0.1.2
 mergedeep==1.3.4
@@ -60,17 +72,23 @@ more-itertools==8.14.0
 multidict==6.0.3
 mypy==1.10.0
 mypy-extensions==1.0.0
+nest-asyncio==1.6.0
 numpy==1.26.4
 OSlash==0.6.3
 packaging==23.1
 pandas==2.2.2
+parso==0.8.4
 pathspec==0.12.1
+pexpect==4.9.0
 pillow==10.3.0
 platformdirs==4.2.2
 plotly==5.22.0
 pluggy==1.0.0
+prompt_toolkit==3.0.48
 protobuf==4.25.3
 psutil==5.9.4
+ptyprocess==0.7.0
+pure_eval==0.2.3
 py==1.11.0
 pyarrow==16.1.0
 pycares==4.3.0
@@ -90,6 +108,7 @@ python-dotenv==1.0.0
 pytz==2024.1
 PyYAML==6.0.1
 pyyaml_env_tag==0.1
+pyzmq==26.2.0
 requests==2.32.3
 rfc3986==1.5.0
 rich==13.7.1
@@ -100,6 +119,7 @@ sniffio==1.3.0
 solana==0.34.0
 solders==0.21.0
 SQLAlchemy==2.0.35
+stack-data==0.6.3
 starlette==0.38.6
 streamlit==1.38.0
 sumtypes==0.1a6
@@ -108,6 +128,8 @@ toml==0.10.2
 tomli==2.0.1
 toolz==0.11.2
 tornado==6.4.1
+tqdm==4.66.5
+traitlets==5.14.3
 types-cachetools==4.2.10
 types-requests==2.31.0.6
 types-urllib3==1.26.25.14
@@ -117,6 +139,7 @@ tzdata==2024.1
 urllib3==1.26.13
 uvicorn==0.31.0
 watchdog==4.0.1
+wcwidth==0.2.13
 websockets==10.4
 wrapt==1.16.0
 yarl==1.8.2

src/page/price_shock.py

@@ -2,61 +2,99 @@
 from asyncio import AbstractEventLoop
 import os
 import time
-from typing import Any
+from typing import Any, TypedDict
 
 from anchorpy import Wallet
 from driftpy.account_subscription_config import AccountSubscriptionConfig
 from driftpy.drift_client import DriftClient
 from driftpy.pickle.vat import Vat
 from lib.api import api
-from lib.page import RPC_STATE_KEY
-from lib.page import VAT_STATE_KEY
-from lib.user_metrics import get_usermap_df
 import pandas as pd
+import plotly.graph_objects as go
 from solana.rpc.async_api import AsyncClient
 import streamlit as st
 
 
-def price_shock_plot(price_scenario_users: list[Any], oracle_distort: float):
-    levs = price_scenario_users
+class UserLeveragesResponse(TypedDict):
+    leverages_none: list[Any]
+    leverages_up: list[Any]
+    leverages_down: list[Any]
+    user_keys: list[str]
+    distorted_oracles: list[str]
+
+
+def create_dataframes(leverages):
+    return [pd.DataFrame(lev) for lev in leverages]
+
+
+def calculate_spot_bankruptcies(df):
+    spot_bankrupt = df[
+        (df["spot_asset"] < df["spot_liability"]) & (df["net_usd_value"] < 0)
+    ]
+    return (spot_bankrupt["spot_liability"] - spot_bankrupt["spot_asset"]).sum()
+
+
+def calculate_total_bankruptcies(df):
+    return -df[df["net_usd_value"] < 0]["net_usd_value"].sum()
+
+
+def generate_oracle_moves(num_scenarios, oracle_distort):
+    return (
+        [-oracle_distort * (i + 1) * 100 for i in range(num_scenarios)]
+        + [0]
+        + [oracle_distort * (i + 1) * 100 for i in range(num_scenarios)]
+    )
+
+
+def price_shock_plot(user_leverages, oracle_distort: float):
+    levs = user_leverages
     dfs = (
-        [pd.DataFrame(levs[2][i]) for i in range(len(levs[2]))]
-        + [pd.DataFrame(levs[0])]
-        + [pd.DataFrame(levs[1][i]) for i in range(len(levs[1]))]
+        create_dataframes(levs["leverages_down"])
+        + [pd.DataFrame(levs["leverages_none"])]
+        + create_dataframes(levs["leverages_up"])
     )
 
-    st.write(dfs)
+    spot_bankruptcies = [calculate_spot_bankruptcies(df) for df in dfs]
+    total_bankruptcies = [calculate_total_bankruptcies(df) for df in dfs]
+
+    num_scenarios = len(levs["leverages_down"])
+    oracle_moves = generate_oracle_moves(num_scenarios, oracle_distort)
+
+    df_plot = pd.DataFrame(
+        {
+            "Oracle Move (%)": oracle_moves,
+            "Total Bankruptcy ($)": total_bankruptcies,
+            "Spot Bankruptcy ($)": spot_bankruptcies,
+        }
+    ).sort_values("Oracle Move (%)")
+
+    df_plot["Perp Bankruptcy ($)"] = (
+        df_plot["Total Bankruptcy ($)"] - df_plot["Spot Bankruptcy ($)"]
+    )
 
-    spot_bankrs = []
-    for df in dfs:
-        spot_b_t1 = df[
-            (df["spot_asset"] < df["spot_liability"]) & (df["net_usd_value"] < 0)
-        ]
-        spot_bankrs.append(
-            (spot_b_t1["spot_liability"] - spot_b_t1["spot_asset"]).sum()
+    fig = go.Figure()
+    for column in [
+        "Total Bankruptcy ($)",
+        "Spot Bankruptcy ($)",
+        "Perp Bankruptcy ($)",
+    ]:
+        fig.add_trace(
+            go.Scatter(
+                x=df_plot["Oracle Move (%)"],
+                y=df_plot[column],
+                mode="lines+markers",
+                name=column,
+            )
         )
 
-    xdf = [
-        [-df[df["net_usd_value"] < 0]["net_usd_value"].sum() for df in dfs],
-        spot_bankrs,
-    ]
-    toplt_fig = pd.DataFrame(
-        xdf,
-        index=["bankruptcy", "spot bankrupt"],
-        columns=[oracle_distort * (i + 1) * -100 for i in range(len(levs[2]))]
-        + [0]
-        + [oracle_distort * (i + 1) * 100 for i in range(len(levs[1]))],
-    ).T
-    toplt_fig["perp bankrupt"] = toplt_fig["bankruptcy"] - toplt_fig["spot bankrupt"]
-    toplt_fig = toplt_fig.sort_index()
-    toplt_fig = toplt_fig.plot()
-
-    toplt_fig.update_layout(
-        title="Bankruptcies in crypto price scenarios",
+    fig.update_layout(
+        title="Bankruptcies in Crypto Price Scenarios",
         xaxis_title="Oracle Move (%)",
         yaxis_title="Bankruptcy ($)",
+        legend_title="Bankruptcy Type",
     )
-    st.plotly_chart(toplt_fig)
+
+    return fig
 
 
 def price_shock_page():
@@ -90,21 +128,26 @@ def price_shock_page():
     )
     st.write(result)
 
-    # price_shock_plot(price_scenario_users, oracle_distort)
+    fig = price_shock_plot(result, oracle_distort)
+    st.plotly_chart(fig)
 
-    # oracle_down_max = pd.DataFrame(price_scenario_users[-1][-1], index=user_keys)
-    # with st.expander(
-    #     str("oracle down max bankrupt count=")
-    #     + str(len(oracle_down_max[oracle_down_max.net_usd_value < 0]))
-    # ):
-    #     st.dataframe(oracle_down_max)
+    oracle_down_max = pd.DataFrame(
+        result["leverages_down"][-1][-1], index=result["user_keys"]
+    )
+    with st.expander(
+        str("oracle down max bankrupt count=")
+        + str(len(oracle_down_max[oracle_down_max.net_usd_value < 0]))
+    ):
+        st.dataframe(oracle_down_max)
 
-    # oracle_up_max = pd.DataFrame(price_scenario_users[1][-1], index=user_keys)
-    # with st.expander(
-    #     str("oracle up max bankrupt count=")
-    #     + str(len(oracle_up_max[oracle_up_max.net_usd_value < 0]))
-    # ):
-    #     st.dataframe(oracle_up_max)
+    oracle_up_max = pd.DataFrame(
+        result["leverages_up"][-1][-1], index=result["user_keys"]
+    )
+    with st.expander(
+        str("oracle up max bankrupt count=")
+        + str(len(oracle_up_max[oracle_up_max.net_usd_value < 0]))
+    ):
+        st.dataframe(oracle_up_max)
 
-    # with st.expander("distorted oracle keys"):
-    #     st.write(distorted_oracles)
+    with st.expander("distorted oracle keys"):
+        st.write(result["distorted_oracles"])