import matplotlib.pyplot as plt
import pandas as pd
from io import StringIO
from itertools import combinations
from math import factorial

# The lesson's toy price model, its 4-district background, and the district x to explain
def toy_model(Z):
    return 100 + 80 * Z['rich'] + 60 * Z['coastal'] + 40 * Z['rich'] * Z['coastal'] + 20 * Z['big']

background = pd.DataFrame({'rich': [0, 1, 0, 0], 'coastal': [0, 0, 1, 0], 'big': [1, 0, 0, 0]})
x = pd.Series({'rich': 1, 'coastal': 1, 'big': 0})
names = list(x.index)
d = len(names)

v = {}
for size in range(d + 1):
    for S in combinations(names, size):
        v[S] = toy_model(background.assign(**{j: x[j] for j in S})).mean()
phi = {j: 0.0 for j in names}
for S, vS in v.items():
    for j in names:
        if j not in S:
            S_j = tuple(k for k in names if k in S or k == j)
            phi[j] += factorial(len(S)) * factorial(d - len(S) - 1) / factorial(d) * (v[S_j] - vS)

# the coalition lattice, left (nobody) to right (everybody)
pos = {(): (0, 0),
       ('rich',): (1, 1.25), ('coastal',): (1, 0), ('big',): (1, -1.25),
       ('rich', 'coastal'): (2, 1.25), ('rich', 'big'): (2, 0), ('coastal', 'big'): (2, -1.25),
       ('rich', 'coastal', 'big'): (3, 0)}

fig, ax = plt.subplots(figsize=(9.5, 4.2))
for S in v:
    for j in names:
        if j in S:
            continue
        S_j = tuple(k for k in names if k in S or k == j)
        (x0, y0), (x1, y1) = pos[S], pos[S_j]
        if j == 'rich':
            w = factorial(len(S)) * factorial(d - len(S) - 1) / factorial(d)
            ax.annotate('', xy=(x1 - 0.2, y1), xytext=(x0 + 0.2, y0),
                        arrowprops=dict(arrowstyle='->', color='#58a6ff', lw=2))
            frac = {1 / 3: '1/3', 1 / 6: '1/6'}[min((1 / 3, 1 / 6), key=lambda f: abs(f - w))]
            ax.text((x0 + x1) / 2 - 0.1, (y0 + y1) / 2 + 0.1, f'+{v[S_j] - v[S]:.0f} × {frac}', ha='right',
                    fontsize=8.5, color='#58a6ff', fontweight='bold',
                    bbox=dict(boxstyle='round,pad=0.15', fc='none', ec='none'))
        else:
            ax.plot([x0 + 0.2, x1 - 0.2], [y0, y1], color='#8b949e', lw=0.8, alpha=0.6)
for S, (px, py) in pos.items():
    label = '{' + ', '.join(S) + '}' if S else '{ } nobody'
    ax.text(px, py, f'{label}\nv = {v[S]:.0f}', ha='center', va='center', fontsize=8.5, color='#1f2328',
            bbox=dict(boxstyle='round,pad=0.35', fc='#dbeaff' if 'rich' in S else '#eceff2',
                      ec='#58a6ff' if 'rich' in S else '#8b949e'))
ax.text(3.45, 1.35, 'Shapley value of rich =\nweighted average of its\n4 blue marginal contributions',
        fontsize=8.5, va='top', color='#8b949e')
ax.text(3.45, 0.45, f"φ(rich) = {phi['rich']:.0f}\nφ(coastal) = {phi['coastal']:.0f}\nφ(big) = {phi['big']:.0f}",
        fontsize=9, va='top', color='#58a6ff')
ax.text(3.45, -0.55, f"sum = {sum(phi.values()):.0f}\n= f(x) − v({{ }})\n= {v[tuple(names)]:.0f} − {v[()]:.0f}",
        fontsize=9, va='top', color='#3fb950')
ax.set_xlim(-0.5, 4.3)
ax.set_ylim(-1.8, 1.8)
ax.axis('off')
ax.set_title('Toy model, x = (rich 1, coastal 1, big 0): v(S) = average prediction with the features in S '
             'set to x', fontsize=9.5, loc='left')

plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()