from pygolem_lite import Shot
import numpy as np

import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
matplotlib.rc('font',  size='10')

import matplotlib.pyplot as plt
from scipy.ndimage import median_filter


def make_img(data, fname, interpolation, start, end):
    plt.figure(figsize=[12, 6])          #size in inches
    plt.imshow(data,
        aspect='auto',
        extent = [start * 1e3, end * 1e3,   #time extents in ms
                   24, 1,],                   #channel idx
        interpolation=interpolation,
        cmap=plt.cm.hot,
        )
    plt.colorbar()
    plt.xlabel('time [ms]')
    plt.ylabel('channel number')
    plt.yticks(np.arange(1, 24+1))
    plt.savefig(fname)
    plt.close()                           #close this figure


def make_plots():
    shot = Shot()
    t, data_za = shot['papouch_za']
    t, data_ko = shot['papouch_ko']
    start = shot['plasma_start']
    end = shot['plasma_end']
    start_idx = np.where(start <= t)[0][0]
    end_idx = np.where(t <= end)[0][-1]
    data = np.hstack([data_za, data_ko]) # NOTE: make sure this is the right order
    data = data[start_idx:end_idx,:]
    data = -data.T                         # transpose to have time as x axis; 
    
    make_img(data, 'bolometers.png', 'none', start, end)
    make_img(median_filter(data, size=(3, 333)), 'bolometers_medfilt-bicube_interp.png', 'bicubic', start, end)


if __name__ == '__main__':
    make_plots()