### import ####################################################################


import os

import numpy as np

import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams['font.size'] = 14
matplotlib.rcParams['savefig.pad_inches'] = 1

import WrightTools as wt


### define ####################################################################


ddir = wt.kit.get_path_matching('Dissertation')

cmap = wt.artists.colormaps['default']
signed_cmap = wt.artists.colormaps['signed']


### TA single vs dual chopping comparison #####################################


if True:
    # get data
    single_p = os.path.join(ddir, 'data', 'method development', 'TA dual chopping', 'SCAN [wmw1, d2] 2016.05.30 11_50_06 single', 'SCAN [wmw1, d2] 2016.05.30 11_50_06 single.data')
    single = wt.data.from_PyCMDS(single_p)
    dual1_p = os.path.join(ddir, 'data', 'method development', 'TA dual chopping', 'SCAN [wmw1, d2] 2016.05.30 12_19_10 dual', 'SCAN [wmw1, d2] 2016.05.30 12_19_10 dual.data')
    dual1 = wt.data.from_PyCMDS(dual1_p)
    dual2_p = os.path.join(ddir, 'data', 'method development', 'TA dual chopping', 'SCAN [wmw1, d2] 2016.05.30 13_32_42 dual', 'SCAN [wmw1, d2] 2016.05.30 13_32_42 dual.data')
    dual2 = wt.data.from_PyCMDS(dual2_p)   
    datas = [single, dual1, dual2]
    pump = wt.units.converter(621.118012, 'nm', 'eV')
    # process data
    for data in datas:
        data.convert('eV')
        data.bring_to_front('signal_diff')
        data.transpose()
    # prepare figure
    fig, gs = wt.artists.create_figure(nrows=3, cols=[1, 'cbar'],
                                       default_aspect=0.5, width='single')
    # plot methods
    def get_vmin_vmax(zi):
        extent = max(zi.max(), -zi.min())
        return -extent, extent
    def plot(data, index, name='name', x=False):
        ax = plt.subplot(gs[index, 0])
        xi = data.axes[1].points
        yi = data.axes[0].points
        zi = data.channels[0].values
        X, Y, Z = wt.artists.pcolor_helper(xi, yi, zi)
        vmin, vmax = get_vmin_vmax(zi)
        ax.pcolor(xi, yi, zi, vmin=vmin, vmax=vmax, cmap=signed_cmap)
        ax.axhline(0, lw=1, c='k')
        ax.set_xlim(xi.min(), xi.max())
        ax.set_ylim(yi.min(), yi.max())
        wt.artists.corner_text(name, ax=ax, fontsize=18)
        ax.grid()
        ax.set_ylabel(single.axes[0].get_label(), fontsize=18)
        ax.axvline(pump, c='k', lw=4, alpha=0.25)
        if x:
            single.axes[1].label_seed = ['probe', 'm']
            ax.set_xlabel(single.axes[1].get_label(), fontsize=18)
            #ax.xaxis.set_label_coords(0.5, 1)
            #ax.xaxis.label._clipon = False
        else:
            plt.setp(ax.get_xticklabels(), visible=False)
    # fill plots
    plot(single, 0, '500 single')
    plot(dual1, 1, '500 dual')
    plot(dual2, 2, '1000 dual', x=True)
    # colorbar
    cax = plt.subplot(gs[:, -1])
    ticks = np.linspace(-1, 1, 9)
    cbar = wt.artists.plot_colorbar(cax, signed_cmap, ticks=ticks,
                                    label='$\mathsf{dI}$')
    # finish
    plt.savefig('TA chopping comparison.png', dpi=300, transparent=True,
                pad_inches=1)
    plt.close(fig)