import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.patches import Patch
from sklearn.model_selection import GroupKFold, KFold, StratifiedKFold, TimeSeriesSplit

# 60 rows as they sit in the file: sorted by class (20% positives at the end)
# and by group (12 patients of unequal size) — which is how real tables often come
n = 60
y = np.array([0] * 48 + [1] * 12)
groups = np.repeat(np.arange(12), [3, 6, 4, 5, 7, 3, 6, 4, 5, 6, 5, 6])
X = np.zeros((n, 1))

schemes = [
    ('KFold (no shuffle)', KFold(5)),
    ('StratifiedKFold', StratifiedKFold(5)),
    ('GroupKFold', GroupKFold(5)),
    ('TimeSeriesSplit', TimeSeriesSplit(5)),
]
ROLE = ['#dddddd', '#58a6ff', '#f0883e']          # not used, train, validation
CLASS = ['#bbbbbb', '#3fb950']
GROUP = ['#8a63d2', '#d2bdf5', '#5a3e8a'] * 4          # alternating shades: one block per patient


def strip(ax, values, colors, y, h=0.7):
    """One horizontal strip, drawn as contiguous runs of equal value."""
    start = 0
    for i in range(1, len(values) + 1):
        if i == len(values) or values[i] != values[start]:
            ax.broken_barh([(start - 0.5, i - start)], (y - h / 2, h), color=colors[values[start]])
            start = i


fig, axes = plt.subplots(2, 2, figsize=(9, 4.6))
for ax, (name, cv) in zip(axes.ravel(), schemes):
    kwargs = {'groups': groups} if isinstance(cv, GroupKFold) else {}
    for i, (tr, va) in enumerate(cv.split(X, y, **kwargs)):
        row = np.zeros(n, dtype=int)
        row[tr], row[va] = 1, 2
        strip(ax, row, ROLE, i + 0.5)
    strip(ax, y, CLASS, 5.7)
    strip(ax, groups, GROUP, 6.5)
    ax.set_yticks([0.5, 1.5, 2.5, 3.5, 4.5, 5.7, 6.5],
                  ['fold 1', 'fold 2', 'fold 3', 'fold 4', 'fold 5', 'class', 'group'], fontsize=7)
    ax.set_ylim(7.0, -0.2)
    ax.set_xlim(-1, n)
    ax.set_xticks([])
    ax.set_title(name, fontsize=9)
    for s in ax.spines.values():
        s.set_visible(False)

axes[1, 0].set_xlabel('rows of the dataset, in file order →', fontsize=8)
axes[1, 1].set_xlabel('time →', fontsize=8)
fig.legend(handles=[Patch(color='#58a6ff', label='train'), Patch(color='#f0883e', label='validation'),
                    Patch(color='#dddddd', label='not used'), Patch(color='#3fb950', label='positive class')],
           loc='lower center', ncol=4, fontsize=8, frameon=False)
plt.tight_layout(rect=(0, 0.06, 1, 1))

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

plt.close()