import sys
import os.path as op
sys.path.append(op.join(op.dirname(__file__), '..'))
from scold import arr_sim, draw, utils

import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib import rc
import numpy as np
import ot
import os.path as op

fig_font_size = 8

plt.rcParams.update({
    "text.usetex": True,
    "font.family": "Helvetica",
    'font.size': fig_font_size
})

rc('text.latex', preamble='\n'.join([
    r'\usepackage{tgheros}',    # helvetica font
    r'\renewcommand\familydefault{\sfdefault} ',
    r'\usepackage[T1]{fontenc}'
]))

# %%
# letters

font = 'arial.ttf'
font_size = 20
scale_mass = True

k = draw.text_array('k', font=font, size=font_size, method=None)
h = draw.text_array('h', font=font, size=font_size, method=None)

# pad and assign to source and target arrays
k_pad, h_pad = utils.pad_for_translation(h.T, k.T, pad=False, constant_values=0.0)

# %%
# plot

fig = plt.figure(layout = 'constrained', figsize = (1.5, 0.8))

gs = GridSpec(1, 3, figure=fig)

axa = fig.add_subplot(gs[0])

pl_rgb = np.zeros((k_pad.shape[0], k_pad.shape[1], 3))
pl_rgb[:, :, 0] = k_pad/k_pad.max()
pl_rgb[:, :, 2] = h_pad/h_pad.max()
axa.imshow(utils.rotate_rgb_hue(1-pl_rgb.transpose((1, 0, 2)), 0.5), interpolation='none')

axa.set_xticks([0,4,8])
axa.set_yticks([0,4,8,12])
axa.set_xticklabels(axa.get_xticks())
axa.set_yticklabels(axa.get_yticks())
axa.spines[['right', 'top']].set_visible(False)

axb = fig.add_subplot(gs[1:])
axb.axis('off')

axb.text(0.1, 1, r'\[\frac{\Sigma\ \ }{\Sigma\ \ }\]', horizontalalignment='left', verticalalignment='top', fontsize=24)

ia_upper = inset_axes(axb, height=0.25, width=1, loc=2, borderpad=0, bbox_to_anchor=(-0.14, 0.027, 1, 1), bbox_transform=axb.transAxes)

ia_upper.axis('off')

intersection = np.transpose(np.minimum(k_pad, h_pad) / k_pad.max())
intersection_pad = np.pad(intersection, pad_width=((1,1), (1,1)), mode='constant', constant_values=0.0)

union = np.transpose(np.maximum(k_pad, h_pad))
union_pad = np.pad(union, pad_width=((1,1), (1,1)), mode='constant', constant_values=0.0)

ia_int_im = ia_upper.imshow(intersection_pad, cmap='gist_gray', interpolation='none', vmin=0.0, vmax=1.0)
# cb = plt.colorbar(ia_int_im)

ia_lower = inset_axes(axb, height=0.25, width=1, loc=3, borderpad=0, bbox_to_anchor=(-0.14, -0.39, 1, 1), bbox_transform=axb.transAxes)

ia_lower.axis('off')

ia_union_im = ia_lower.imshow(union_pad, cmap='gist_gray', interpolation='none', vmin=0.0, vmax=1.0)

fig.savefig(op.join('..', 'fig', 'intro_jaccard_examples.pdf'))
fig.savefig(op.join('..', 'fig', 'intro_jaccard_examples.png'))
fig.savefig(op.join('..', 'fig', 'intro_jaccard_examples.svg'))

print(f'Jaccard = {np.sum(intersection) / np.sum(union)}')