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


import os

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

import WrightTools as wt


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


gd_path = wt.kit.get_path_matching('Google Drive')
directory = os.path.dirname(__file__)

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


### process ###################################################################


if True:
    fig, gs = wt.artists.create_figure(nrows=3, cols=[1, 1, 'cbar'],
                                       default_aspect=0.5, width='double')
    def get_vmin_vmax(zi):
        extent = max(zi.max(), -zi.min())
        return -extent, extent
    def plot(data, col_index, name='name'):    
        # raw
        ax = plt.subplot(gs[0, col_index])
        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=cmap)
        ax.axhline(0, lw=1, c='k')
        ax.set_xlim(xi.min(), xi.max())
        ax.set_ylim(yi.min(), yi.max())
        ax.set_title(name, fontsize=18)
        wt.artists.corner_text('raw', ax=ax, fontsize=18)
        # leveled
        ax = plt.subplot(gs[1, col_index])
        data.level(0, 'd2', -3)
        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)
        ax.pcolor(xi, yi, zi, vmin=vmin, vmax=vmax, cmap=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('leveled', ax=ax, fontsize=18)
        # smoothed
        ax = plt.subplot(gs[2, col_index])
        data.smooth(2)
        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)
        ax.pcolor(xi, yi, zi, vmin=vmin, vmax=vmax, cmap=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('smoothed', ax=ax, fontsize=18)
    # single
    p = os.path.join(directory, 'SCAN [wmw1, d2] 2016.05.30 11_50_06 single', 'SCAN [wmw1, d2] 2016.05.30 11_50_06 single.data')
    data = wt.data.from_PyCMDS(p)
    data.transpose()
    data.bring_to_front('signal_diff')
    data.channels[0].values *= -1
    plot(data, 0, name='single')
    # dual
    p = os.path.join(directory, 'SCAN [wmw1, d2] 2016.05.30 13_32_42 dual', 'SCAN [wmw1, d2] 2016.05.30 13_32_42 dual.data')
    data = wt.data.from_PyCMDS(p)
    data.transpose()
    data.bring_to_front('signal_diff')
    data.channels[0].values *= -1
    plot(data, 1, name='dual')
    # colorbar
    cax = plt.subplot(gs[:, -1])
    matplotlib.colorbar.ColorbarBase(cax, cmap=cmap, ticks=[-1, 0, 1])
    # finish
    plt.savefig('compare.png', dpi=300, transparent=True, pad_inches=1)
    plt.close('fig')