import os.path, glob
from os.path import join
import json

import cmcrameri.cm
import h5py
import pandas as pd

import numpy as np
import matplotlib.pyplot as pl

from helper import *
from style import *

# path data repo
path_data_repo = "../data_supplementary"

# timestemps of the images, they are also the key of the h5 file as "timestemp_{time}"
times = [13.51,13.68,14.20,15.05,16.25,17.28,17.79,18.65]

set_plot_style()
rows = 2
cols = int(len(times))-4
fig, ax  = pl.subplots(nrows=rows,ncols=cols ,figsize=(8,4), sharex=True, sharey=True, layout="constrained")
fontsize = 12

cmap = cmcrameri.cm.roma_r


# open dataframe dataframe
h5f = h5py.File(os.path.join(path_data_repo, fr'vortex_ring_decay.h5'), 'r')
# extract additional information in attributes
discretization_length = h5f.attrs["discretization_length"]  # µm
zSize = h5f.attrs["zSize"]
ySize = h5f.attrs["ySize"]
f = h5f.attrs["f"]
I_m = h5f.attrs["I_m"]
I_dc = h5f.attrs["I"]


v_max = 200
for i, file in enumerate(times):
    # get data
    key = f"timestemp_{times[i]}"
    print(key)
    data = h5f[key]

    # rescale data to 1/µm
    data = np.asarray(data) * 1/discretization_length**3

    x_max = data.shape[1] * discretization_length
    y_max = data.shape[0] * discretization_length
    extent = (-x_max/2, x_max/2, -y_max /2 ,y_max/2)

    #print()
    im = ax[int(i/4),i%4].imshow(data,aspect='auto', origin='lower', extent=extent, vmax = v_max,cmap=cmap,interpolation="bilinear")

    ax[int(i/4),i%4].text(x=1.7, y=2.7, s=f"{times[i]:.2f} ms", ha='center', verticalalignment='center', fontsize=8,
                             bbox={'boxstyle': 'round', 'fc': RPTU_COLORS["mango25pct"], 'ec': RPTU_COLORS["mango"],
                                   'ls': '-', 'lw': 1.5}, zorder=5)


h5f.close()


ticks = [-2, 0, 2]

for i in range(4):
    ax[1,i].set_xlabel('y [µm]',fontsize = fontsize)
    ax[1, i].set_xticks(ticks)
    ax[1, i].set_xticklabels(ticks, fontsize=fontsize)


ax[0,0].set_ylabel('z [µm]',fontsize = fontsize)
ax[1,0].set_ylabel('z [µm]',fontsize = fontsize)


ax[0, 0].set_yticks(ticks)
ax[0, 0].set_yticklabels(ticks, fontsize=fontsize)
ax[1, 0].set_yticks(ticks)
ax[1, 0].set_yticklabels(ticks, fontsize=fontsize)



cbar = fig.colorbar(im, ax = ax[:, 3], shrink = 1)
cbar.ax.set_ylabel(r'$ n$ [1/µm³]',fontsize =fontsize )
ticks = [200,100,0]
cbar.ax.set_yticks(ticks)
cbar.ax.set_yticklabels(ticks, fontsize = fontsize)

#pl.savefig("vortex_ring.pdf")
pl.show()