import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.io as pio

pio.renderers.default = "plotly_mimetype+notebook_connected"
GOLD, SLATE, GREY, INK = "#D99A1E", "#4F6D9A", "#8A8F98", "#4A4D55"

def toy_economy(spread, n=10_000, seed=42):
    """Log-normal incomes. A bigger spread means a longer tail and more inequality."""
    rng = np.random.default_rng(seed)
    return rng.lognormal(mean=10.8, sigma=spread, size=n)

def lorenz(incomes):
    """Cumulative share of people (poorest first) against their cumulative share of income."""
    x = np.sort(incomes)
    people = np.arange(len(x) + 1) / len(x)
    income = np.concatenate([[0], np.cumsum(x) / x.sum()])
    return people, income

def gini(incomes):
    """One minus twice the area under the Lorenz curve."""
    people, income = lorenz(incomes)
    return 1 - 2 * np.trapezoid(income, people)

spreads = np.round(np.arange(0.2, 1.21, 0.1), 1)
grid = np.linspace(0, 1, 101)

fig = go.Figure()
fig.add_trace(go.Scatter(x=grid, y=grid, mode="lines", name="Perfect equality",
                         line=dict(color=GREY, dash="dash", width=1.5), hoverinfo="skip"))
for i, s in enumerate(spreads):
    people, income = lorenz(toy_economy(s))
    curve = np.interp(grid, people, income)
    fig.add_trace(go.Scatter(
        x=grid, y=curve, mode="lines", name=f"spread {s}", visible=(s == 0.6),
        line=dict(color=GOLD, width=3), fill="tonexty", fillcolor="rgba(217,154,30,0.18)",
        hovertemplate="The poorest %{x:.0%} earn %{y:.1%} of all income<extra></extra>"))

steps = []
for i, s in enumerate(spreads):
    g = gini(toy_economy(s))
    visible = [True] + [j == i for j in range(len(spreads))]
    steps.append(dict(method="update", label=f"{s}",
                      args=[{"visible": visible}, {"title.text": f"Spread {s}: Gini = {g:.2f}"}]))

start = int(np.where(spreads == 0.6)[0][0])
fig.update_layout(
    title=dict(text=steps[start]["args"][1]["title.text"], x=0, font=dict(size=16)),
    sliders=[dict(active=start, steps=steps, currentvalue=dict(prefix="Spread: "), pad=dict(t=40))],
    xaxis=dict(title="Share of people, poorest first", tickformat=".0%", range=[0, 1]),
    yaxis=dict(title="Share of income", tickformat=".0%", range=[0, 1]),
    showlegend=False, template="simple_white", height=520, margin=dict(l=60, r=20, t=60, b=40),
    font=dict(family="IBM Plex Sans, sans-serif"),
)
fig.show()

table = pd.DataFrame({"spread": spreads, "gini": [gini(toy_economy(s)) for s in spreads]})

fig2 = go.Figure(go.Scatter(
    x=table["spread"], y=table["gini"], mode="lines+markers",
    line=dict(color=GOLD, width=3), marker=dict(size=8, color=GOLD),
    hovertemplate="Spread %{x}<br>Gini %{y:.2f}<extra></extra>"))
fig2.update_layout(
    xaxis=dict(title="Spread of the toy economy"), yaxis=dict(title="Gini coefficient", range=[0, 0.65]),
    template="simple_white", height=380, margin=dict(l=60, r=20, t=20, b=40),
    font=dict(family="IBM Plex Sans, sans-serif"),
)
fig2.show()

def quintile_shares(incomes):
    x = np.sort(incomes)
    return [chunk.sum() / x.sum() for chunk in np.array_split(x, 5)]

labels = ["Poorest 20%", "Second", "Middle", "Fourth", "Richest 20%"]
fig3 = go.Figure()
for s, color in [(0.3, GREY), (0.6, SLATE), (1.0, GOLD)]:
    fig3.add_trace(go.Bar(x=labels, y=quintile_shares(toy_economy(s)), name=f"Spread {s}",
                          marker_color=color, hovertemplate="%{x}: %{y:.1%} of income<extra></extra>"))
fig3.update_layout(barmode="group", yaxis=dict(title="Share of income", tickformat=".0%"),
                   template="simple_white", height=400, margin=dict(l=60, r=20, t=20, b=40),
                   legend=dict(orientation="h", y=1.08), font=dict(family="IBM Plex Sans, sans-serif"))
fig3.show()