import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
from pathlib import Path

root_dir = Path('/media/akovalev/CrucialX10/work/umh/data/raw/export/things_mua_export')

train_table = pd.read_csv(root_dir / 'train_table.csv')
test_table = pd.read_csv(root_dir / 'test_table.csv')

mua_train = np.load(root_dir / 'train_mua.npy', mmap_mode='r')
mua_test = np.load(root_dir / 'test_mua.npy', mmap_mode='r')

stats = np.load(root_dir / 'global_stats.npz')
means, stds = stats['means'], stats['stds']
eps = 1e-6

print("Train dataset size: ", mua_train.shape)
print("Test dataset size: ", mua_test.shape)
print(f'Shape: [N, C, T] where \nN is number of samples, \nC is number of channels, \nT is number of time points')



def plot_pair(mua, image, path_to_save=None):
    fig, axs = plt.subplots(1, 2, figsize=(10, 5))
    axs[0].imshow(mua, aspect='auto', cmap='seismic', vmin=-5, vmax=5)
    axs[1].imshow(image, aspect='auto')
    if path_to_save is not None:
        plt.savefig(path_to_save)
    plt.close()

def preproc_mua(mua, means, stds, eps):
    mua = (mua - means) / stds + eps
    mua = np.sign(mua) * np.log(1 + np.abs(mua))
    return mua


# Show samples for train set.
test_idxs = [0, 5000, 10000, 15000, 20000]

for idx in test_idxs:
    mua = mua_train[idx]
    mua = preproc_mua(mua, means, stds, eps)

    image_relative_path = train_table[train_table.idx == idx]['image_path'].values[0]
    image_path = root_dir / 'images' / image_relative_path
    image = plt.imread(image_path)

    plot_pair(mua, image, path_to_save=f'train_sample_{idx}.png')


test_idxs = [0, 10, 20, 100, 200]

for idx in test_idxs:
    mua = mua_test[idx]
    mua = preproc_mua(mua, means, stds, eps)

    image_relative_path = test_table[test_table.idx == idx]['image_path'].values[0]
    image_path = root_dir / 'images' / image_relative_path
    image = plt.imread(image_path)

    plot_pair(mua, image, path_to_save=f'test_sample_{idx}.png')