import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import ListedColormap
from scipy.cluster.hierarchy import dendrogram, leaves_list, linkage

# A synthetic gene-expression table: 40 genes x 12 samples (6 tumour, 6 normal).
# Genes 0-14 go up in tumours, 15-29 go down, 30-39 are noise; every gene has
# its own baseline level and amplitude. Rows and columns arrive shuffled.
rng = np.random.default_rng(3)
group = np.r_[np.ones(6), -np.ones(6)]                         # +1 tumour, -1 normal
module = np.r_[np.zeros(15, int), np.ones(15, int), np.full(10, 2)]
effect = np.array([1.0, -1.0, 0.0])[module]
base, amp = rng.uniform(2, 10, 40), rng.uniform(0.8, 2.0, 40)
E = base[:, None] + amp[:, None] * effect[:, None] * group[None, :] + rng.normal(0, 0.6, (40, 12))
samples = np.array([f'T{i}' for i in range(1, 7)] + [f'N{i}' for i in range(1, 7)])
r0, c0 = rng.permutation(40), rng.permutation(12)
E, module, samples = E[r0][:, c0], module[r0], samples[c0]

# What a clustered heat map does: z-score each gene, cluster rows and columns
# (average linkage, correlation distance), and draw both in leaf order.
Ez = (E - E.mean(axis=1, keepdims=True)) / E.std(axis=1, keepdims=True)
Zr = linkage(Ez, 'average', metric='correlation')
Zc = linkage(Ez.T, 'average', metric='correlation')
rows, cols = leaves_list(Zr), leaves_list(Zc)

GRAY = '#8b949e'
BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})
MOD = ListedColormap([ORANGE, BLUE, GRAY])
fig = plt.figure(figsize=(10, 5.6))
gs = fig.add_gridspec(2, 6, width_ratios=[4, 0.25, 1.2, 1.4, 0.25, 4],
                      height_ratios=[1, 5], wspace=0.05, hspace=0.04)

def heat(ax, M, labels):
    im = ax.imshow(M, aspect='auto', cmap='RdBu_r', vmin=-2.2, vmax=2.2, interpolation='nearest')
    ax.set_xticks(range(len(labels)), labels, rotation=90, fontsize=8)
    ax.set_yticks([])
    return im

# left: as measured
ax = fig.add_subplot(gs[1, 0])
heat(ax, Ez, samples)
ax.set_title('as measured (each gene z-scored)', fontsize=9.5, loc='left')
ax.set_ylabel('40 genes', fontsize=9)
ax = fig.add_subplot(gs[1, 1])
ax.imshow(module[:, None], aspect='auto', cmap=MOD, vmin=0, vmax=2, interpolation='nearest')
ax.set_xticks([0], ['true\nmodule'], fontsize=7.5)
ax.set_yticks([])

# right: reordered, with both trees
ax = fig.add_subplot(gs[1, 3])
dendrogram(Zr, orientation='left', ax=ax, no_labels=True, link_color_func=lambda k: GRAY)
ax.invert_yaxis()
ax.axis('off')
ax = fig.add_subplot(gs[1, 4])
ax.imshow(module[rows][:, None], aspect='auto', cmap=MOD, vmin=0, vmax=2, interpolation='nearest')
ax.set_xticks([0], ['true\nmodule'], fontsize=7.5)
ax.set_yticks([])
axh = fig.add_subplot(gs[1, 5])
im = heat(axh, Ez[rows][:, cols], samples[cols])
axt = fig.add_subplot(gs[0, 5])
dendrogram(Zc, ax=axt, no_labels=True, link_color_func=lambda k: GRAY)
axt.axis('off')
axt.set_title('the same numbers, reordered by the two trees', fontsize=9.5, loc='left')

cax = fig.add_axes([0.72, 0.035, 0.18, 0.018])
cb = fig.colorbar(im, cax=cax, orientation='horizontal')
cb.ax.tick_params(labelsize=7.5)
fig.text(0.71, 0.035, 'z-score ', fontsize=8.5, ha='right', va='bottom')
fig.text(0.02, 0.02, 'module:', fontsize=8.5)
for i, (name, c) in enumerate([('up in tumour', ORANGE), ('down in tumour', BLUE), ('noise', GRAY)]):
    fig.text(0.08 + i * 0.12, 0.02, f'■ {name}', color=c, fontsize=8.5)

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

plt.close()