# file to evaluate a measurement
import os.path, glob
from os.path import join
import json

import matplotlib.pyplot as plt
import pandas as pd
from scipy.ndimage import gaussian_filter

import numpy as np
import matplotlib.pyplot as pl
from scipy.optimize import curve_fit

from helper import *
from style import *
import h5py

path_data_repo = "../data_supplementary"

meas_list = [{"path": "speed_of_sound.h5"}]

colormap = cmaps.vik
v_max = 40

for meas in meas_list:
    path = join(path_data_repo, meas["path"])

    h5_file = h5py.File(path, 'r')
    image = h5_file['image']

    pixel_size = h5_file.attrs['pixel_size']
    magnification = h5_file.attrs['magnification']
    conversion_to_um = h5_file.attrs['conversion_to_um']
    t_max = h5_file.attrs['t_max']

    x_max = image.shape[1] * conversion_to_um

    # transform the atomic density back to atom space
    image = np.asarray(image) * pixel_size / magnification # [N/µm] ; Pixel size = 6.45 µm/Pixel

    fig, ax = pl.subplots(nrows=1, ncols=1, dpi=300, figsize=(3 * 1.61, 3))

    im = ax.imshow(image, cmap=colormap, aspect='auto', origin='lower',
                   extent=(0, x_max, 0, t_max), vmin=-v_max, vmax=v_max)

    cbar = pl.colorbar(im, ax=ax)
    cbar.ax.set_ylabel(r'$\Delta n $ [1/µm]')
    ax.set_ylabel("Time [ms]")
    ax.set_xlabel('Distance [µm]')

    fit = True
    if fit:
        min_val = 0
        max_val = t_max
        smoothed = gaussian_filter(image, sigma=3)

        velocities = []
        velocities_err = []
        for mask in [(0, 200), (200, -1)]:
            offset = mask[0]
            sign = np.sign(mask[1])
            mask = slice(*mask)
            #print(mask)
            x_vals = []
            y_vals = []
            for i in range(smoothed.shape[0] - 20):
                x_vals.append(conversion_to_um * (np.argmin(smoothed[i, mask]) + offset))
                y_vals.append(i / smoothed.shape[0] * (max_val - min_val))
            x_vals = np.asarray(x_vals)
            y_vals = np.asarray(y_vals)

            def linear(x, m, b):
                return x / m + b


            mask = y_vals < 10
            p_opt, p_cov = curve_fit(linear, x_vals[mask], y_vals[mask], p0=[-1.5 * sign, -25])
            p_err = np.sqrt(np.diag(p_cov))
            plt.errorbar(x_vals[mask], y_vals[mask], **get_style(color='orange'), markersize=5)
            plt.errorbar(x_vals, linear(x_vals, *p_opt), color='forestgreen', lw=2)
            velocities.append(round(p_opt[0], 2))
            velocities_err.append(round(abs(p_err[0]), 2))


    fig.tight_layout()
    # pl.savefig(f'speed_of_sound.png'))
    pl.show()