# %%
import matplotlib.pyplot as plt
import numpy as np

rng = np.random.default_rng(0)
n = 30
y_true = rng.uniform(low=0, high=10, size=n)
y_pred = y_true * rng.uniform(low=0.9, high=1.1, size=n)
exposure = rng.integers(low=0, high=10, size=n)


def lorenz_curve_linspace(frequency, exposure, label=None, use_cumulated_exposure=True):
    ranking = np.argsort(frequency)
    ranked_frequencies = frequency[ranking]
    ranked_exposure = exposure[ranking]
    cumulated_claims = np.cumsum(ranked_frequencies * ranked_exposure)
    cumulated_claims = cumulated_claims / cumulated_claims[-1]

    if use_cumulated_exposure:
        cumulated_exposure = np.cumsum(ranked_exposure).astype(np.float64)
        cumulated_exposure /= cumulated_exposure[-1]
    else:
        cumulated_exposure = np.linspace(0, 1, len(frequency))

    plt.scatter(
        cumulated_exposure,
        cumulated_claims,
        marker=".",
        alpha=0.5,
        label=label,
    )
    return cumulated_exposure, cumulated_claims


y_true_repeated = y_true.repeat(exposure)
y_pred_repeated = y_pred.repeat(exposure)
sample_weight_repeated = np.ones_like(y_pred_repeated)
res_repeated = lorenz_curve_linspace(
    y_pred_repeated, sample_weight_repeated, label="repeated"
)

y_true_weighted = y_true
y_pred_weighted = y_pred
sample_weight_weighted = exposure
res_weighted = lorenz_curve_linspace(
    y_pred_weighted, sample_weight_weighted, label="weighted"
)

plt.legend()
plt.title("Lorenz curve using cumulated exposure as x-axis");
# %%
res_repeated = lorenz_curve_linspace(
    y_pred_repeated,
    sample_weight_repeated,
    label="repeated",
    use_cumulated_exposure=False,
)
res_weighted = lorenz_curve_linspace(
    y_pred_weighted,
    sample_weight_weighted,
    label="weighted",
    use_cumulated_exposure=False,
)

plt.legend()
plt.title("Lorenz curve using linear exposure as x-axis");