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


import os
import collections

import numpy as np

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

import WrightTools as wt
import WrightData as wd


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


width = 14

gd_path = wt.kit.get_path_matching('Google Drive')

data_path = os.path.join(gd_path, 'Blaise', 'unpublished Thompson - Automated OPA Tuning')


### recent examples ###########################################################


if False:
    # prepare figure
    wf = 0.01
    cols = ['cbar', 1, 1, 'cbar']
    aspects = [[[0, 1], 0.5]]
    fig, gs = wt.artists.create_figure(width=width/2., nrows=2, cols=cols, hspace=0.5, wspace=0.5, aspects=aspects)
    # dOD colorbar
    cax = plt.subplot(gs[:, 0])
    limit = 0.05
    ticks = np.linspace(-limit, limit, 5)
    wt.artists.plot_colorbar(cax=cax, cmap='signed', label='dOD', ticklocation='left', ticks=ticks)
    # example 1 (MoS2 TA)
    ax = plt.subplot(gs[0, 1])
    p = os.path.join(gd_path, 'MX2', 'CMDS', '2016.05.26', 'SCAN [w2, wmw1, d2] 2016.05.30 18_17_55 TA', 'SCAN [w2, wmw1, d2] 2016.05.30 18_17_55 TA.data')
    data = wt.data.from_PyCMDS(p, verbose=False)
    data = data.chop('w2', 'wmw1', verbose=False)[11]
    data = data.split('wmw1', 13250, verbose=False)[1]
    data.smooth(2)
    data.dOD('dI', 'I')
    data.dI.clip(-limit, limit)
    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(X, Y, Z, vmin=-limit, vmax=limit, cmap=wt.artists.colormaps['signed'])
    wt.artists.diagonal_line(xi, yi)
    ax.grid()
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_xlabel('$\\mathsf{\\bar\\nu_{probe}=\\bar\\nu_{m}}$', fontsize=18)
    ax.set_ylabel('$\\mathsf{\\bar\\nu_{pump}}$', fontsize=18)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_ylim(yi.min(), yi.max())
    # example 2 (MoS2 TrEE)
    ax = plt.subplot(gs[0, 2])
    data = wd.get('2015.12 Czech', check_remote=False)['movie']
    data = data.chop('w1', 'd2')[29]
    data.transpose()
    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(X, Y, Z, cmap=wt.artists.colormaps['default'])
    ax.grid()
    ax.axhline(0, c='k', lw=1)
    ax.axvline(data.w2.points, c='k', lw=4, alpha=0.25)
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_xlabel('$\\mathsf{\\bar\\nu_{1}=\\bar\\nu_{m}}$', fontsize=18)
    ax.set_ylabel('$\\mathsf{\\tau_{21}}$', fontsize=18)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_ylim(yi.min(), yi.max())
    # example 3 (PbSe TrEE)
    ax = plt.subplot(gs[1, 1])
    data = wt.data.from_pickle('PbSe.p')
    data = data.chop('w2', 'w1')[12]
    xi = data.axes[1].points
    yi = data.axes[0].points
    zi = data.channels[0].values
    zi -= zi.min()
    zi += 0.001
    zi /= zi.max()
    zi **= 0.5
    X, Y, Z = wt.artists.pcolor_helper(xi, yi, zi)
    ax.pcolor(X, Y, Z, cmap=wt.artists.colormaps['default'])
    wt.artists.diagonal_line(xi, yi)
    ax.grid()
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_xlabel('$\\mathsf{\\bar\\nu_{1}=\\bar\\nu_{m}}$', fontsize=18)
    ax.set_ylabel('$\\mathsf{\\bar\\nu_{2}}$', fontsize=18)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_ylim(yi.min(), yi.max())
    # example 4 (TRSF)
    ax = plt.subplot(gs[1, 2])
    data = wd.get('2013.10 Boyle', check_remote=False)['full 2D TRSF']
    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(X, Y, Z, cmap=wt.artists.colormaps['default'])
    wt.artists.diagonal_line(xi, yi)
    ax.grid()
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_xlabel('$\\mathsf{\\bar\\nu_{1}}$', fontsize=18)
    ax.set_ylabel('$\\mathsf{\\bar\\nu_{2}}$', fontsize=18)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_ylim(yi.min(), yi.max())
    # homodyne colorbar
    cax = plt.subplot(gs[:, -1])
    wt.artists.plot_colorbar(cax=cax, label='amplitude')
    # finish
    plt.savefig('frequency domain examples.png', dpi=300, transparent=True, pad_inches=1)
    plt.close(fig)

    
### curve #####################################################################


if False:
    # prepare figure
    fig, gs = wt.artists.create_figure(width=width/2., nrows=2, cols=[1],
                                       default_aspect=0.3)
    # get curve
    d = os.path.join(gd_path, 'Development', 'OPA tuning', 'tuning records', 'TOPAS-C 10743', 'signal', '2016.05.28', 'OPA1 (10743) curves')
    curves = [os.path.join(d, 'OPA1 (10743) base - 2016.05.28 12_56_14.crv'),
              os.path.join(d, 'OPA1 (10743) mixer1 - 2015.11.09'),
              os.path.join(d, 'OPA1 (10743) mixer2 - 2016.05.15 10_23_48'),
              os.path.join(d, 'OPA1 (10743) mixer3 - 2013.06.01'),]
    curve = wt.tuning.curve.from_TOPAS_crvs(curves, 'TOPAS-C', 'NON-NON-NON-Sig')
    # preamp
    ax = plt.subplot(gs[0, 0])
    plt.setp(ax.get_xticklabels(), visible=False)
    ax.grid()
    ax.set_xticks([1200, 1400, 1600])
    # C1
    ax.scatter(curve.colors, curve.motors[0].positions, c='b', edgecolor='none')
    ax.plot(curve.colors, curve.motors[0].positions, c='b', lw=2)
    gone = [0, 7]
    for i, tl in enumerate(ax.get_yticklabels()):
        tl.set_color('b')
        if i in gone:
            tl.set_visible(False)
    ax.set_ylabel('C1 (deg)', color='b', fontsize=18)
    # D1
    ax = ax.twinx()
    ax.scatter(curve.colors, curve.motors[1].positions, c='r', edgecolor='none')
    ax.plot(curve.colors, curve.motors[1].positions, c='r', lw=2)
    gone = [5]
    for i, tl in enumerate(ax.get_yticklabels()):
        tl.set_color('r')
        if i in gone:
            tl.set_visible(False)
    ax.set_ylabel('D1 (mm)', color='r', fontsize=18)
    # poweramp
    ax = plt.subplot(gs[1, 0])
    ax.grid()
    xi = np.linspace(1100, 1700, 4)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_xticks([1200, 1400, 1600])
    ax.set_xlabel('setpoint (nm)', fontsize=18)
    # C2
    ax.scatter(curve.colors, curve.motors[2].positions, c='b', edgecolor='none')
    ax.plot(curve.colors, curve.motors[2].positions, c='b', lw=2)
    gone = [0, 9]
    for i, tl in enumerate(ax.get_yticklabels()):
        tl.set_color('b')
        if i in gone:
            tl.set_visible(False)
    ax.set_ylabel('C2 (deg)', color='b', fontsize=18)
    # D2
    ax = ax.twinx()
    ax.scatter(curve.colors, curve.motors[3].positions, c='r', edgecolor='none')
    ax.plot(curve.colors, curve.motors[3].positions, c='r', lw=2)
    gone = [0, 5]
    for i, tl in enumerate(ax.get_yticklabels()):
        tl.set_color('r')
        if i in gone:
            tl.set_visible(False)
    ax.set_ylabel('D2 (deg)', color='r', fontsize=18)
    yticks = ax.yaxis.get_major_ticks()
    # finish
    plt.savefig('curve.png', dpi=300, transparent=True, pad_inches=1)
    plt.close(fig)


### tuning range ##############################################################


if False:
    # prepare figure
    fig, gs = wt.artists.create_figure(width=width, nrows=1, cols=[1],
                                       default_aspect=0.15)
    # populate
    ax = plt.subplot(gs[0, 0])
    cs = ['k', 'grey', 'orange', 'r', 'c', 'm', 'g', 'b', 'k']
    ranges = collections.OrderedDict()
    ranges['Idler'] = [2600, 1600]
    ranges['Signal'] = [1600, 1150]
    ranges['SHI'] = [800, 1200]
    ranges['SHS'] = [800, 580]
    ranges['SFI'] = [600, 533]
    ranges['SFS'] = [533, 480]
    ranges['4HI'] = [590, 400]
    ranges['4HS'] = [400, 290]
    height = 0
    for name, limits in ranges.items():
        ax.plot(limits, [height, height], lw=10, label=name, c=cs[height+1])
        height += 1
    labels = [''] + ranges.keys() + ['']
    [i.set_color(c) for i, c in zip(plt.gca().get_yticklabels(), cs)]
    ax.set_yticklabels(labels)
    ax.set_xlim(290, 2000)
    ax.set_ylim(-1, len(ranges))
    ax.set_xlabel('OPA output (nm)', fontsize=18)
    ax.grid()
    # finish
    plt.savefig('ranges.png', dpi=300, transparent=True, pad_inches=1)
    plt.close(fig)


### huge preamp ###############################################################


if False:
    # prepare figure
    fig, gs = wt.artists.create_figure(width=width, nrows=1, hspace=1,
                                       cols=[ 1, 'cbar', 0.1, 1, 'cbar'],
                                       aspects=[[[0, 0], 1]])
    # get data
    p = 'TOPAS_C_full_preamp.p'
    data = wt.data.from_pickle(p)
    #data.transpose()
    # clip based on intensity
    data.amplitude.clip(zmin=0.1, zmax=5)
    # clip based on width
    data.width.clip(zmin=10, zmax=100)
    # clip based on center
    data.center.clip(zmin=1140, zmax=1620)
    # share NaNs
    data.share_nans()
    data.heal('center')
    data.heal('amplitude')
    # intensity
    ax = plt.subplot(gs[0, 0])
    cax = plt.subplot(gs[0, 1])
    xi = data.c1.points
    yi = data.d1.points
    zi = data.amplitude.values
    zi /= np.nanmax(zi)
    X, Y, Z = wt.artists.pcolor_helper(xi, yi, zi)
    cmap=wt.artists.colormaps['default']
    cmap.set_under([0.75]*3, 1)
    mappable = ax.pcolor(X, Y, Z, vmin=0, vmax=1, cmap=cmap)
    plt.colorbar(mappable=mappable, cax=cax, ticklocation='right')
    cax.set_ylabel('intensity', fontsize=18)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_ylim(1.35, 1.8)
    ax.contour(xi, yi, zi, 5, colors='k')
    ax.set_xlabel('C1 (deg)', fontsize=18)
    ax.set_ylabel('D1 (mm)', fontsize=18)
    ax.grid()
    # color
    ax = plt.subplot(gs[0, 3])
    cax = plt.subplot(gs[0, 4])
    xi = data.c1.points
    yi = data.d1.points
    zi = data.center.values
    X, Y, Z = wt.artists.pcolor_helper(xi, yi, zi)
    cmap=wt.artists.colormaps['rainbow']
    cmap.set_under([0.75]*3, 1)
    mappable = ax.pcolor(X, Y, Z, vmin=np.nanmin(Z), vmax=np.nanmax(Z), cmap=cmap)
    plt.colorbar(mappable=mappable, cax=cax)
    cax.set_ylabel('color (nm)', fontsize=18)
    ax.set_xlim(xi.min(), xi.max())
    ax.set_ylim(1.35, 1.8)   
    ax.contour(xi, yi, zi, 25, colors='k')
    ax.set_xlabel('C1 (deg)', fontsize=18)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.grid()
    # finish
    plt.savefig('preamp.png', dpi=300, transparent=True, pad_inches=1)
    plt.close(fig)


### actual preamp #############################################################


if False:
    p = os.path.join(gd_path, 'Development', 'OPA tuning', 'tuning records', 'TOPAS-C 10743', 'signal', '2016.05.14', 'TOPAS-C AUTOTUNE [w1, w1_Delay_1, wa] 2016.05.14 20_25_26', 'TOPAS-C AUTOTUNE [w1, w1_Delay_1, wa] 2016.05.14 20_25_26.data')
    d = r'C:\Users\blais\Google Drive\Development\OPA tuning\tuning records\TOPAS-C 10743\signal\2016.04.25\OPA1 (10743) curves'    
    curves = [os.path.join(d, 'OPA1 (10743) base - 2016.04.25 22_46_02.crv'),
              os.path.join(d, 'OPA1 (10743) mixer1 - 2015.11.09'),
              os.path.join(d, 'OPA1 (10743) mixer2 - 2016.03.22 23_31_12'),
              os.path.join(d, 'OPA1 (10743) mixer3 - 2013.06.01'),]
    wt.tuning.TOPAS_C.process_preamp_motortune(1, p, curves, save=True)


### huge poweramp #############################################################


if False:
    # prepare figure
    fig, gs = wt.artists.create_figure(width=width, nrows=1, hspace=1,
                                       cols=[ 1, 'cbar', 0.1, 1, 'cbar'],
                                       aspects=[[[0, 0], 4/3.]])
    # get data
    p = 'TOPAS_C_full_poweramp_moments.p'
    data = wt.data.from_pickle(p, verbose=False)
    setpoints = np.linspace(1160, 1600, 12)
    # clip, share NaNs
    data.integral.clip(zmin=0.1, zmax=500)
    data.one.clip(zmin=1140, zmax=1620)
    # plot method
    def plot(subplot_spec, cax, channel_index, cmap, yticklabels=False):
        inner_gs = matplotlib.gridspec.GridSpecFromSubplotSpec(4, 3, subplot_spec=subplot_spec, wspace=0.0, hspace=0.0)
        # fill in
        for i, g in enumerate(inner_gs):
            ax = plt.subplot(g)
            if not ax.is_last_row():
                plt.setp(ax.get_xticklabels(), visible=False)
            if not ax.is_first_col():
                plt.setp(ax.get_yticklabels(), visible=False)
            if not yticklabels:
                plt.setp(ax.get_yticklabels(), visible=False)
            ax.grid()
            if i == 10:
                ax.set_xlabel('$\mathsf{\Delta}$C2 (deg)', fontsize=18)
            cm = wt.artists.colormaps[cmap]
            cm.set_under([0.75]*3)
            setpoint = setpoints[i]
            d = data.chop('d2', 'c2', {'w1': [setpoint, 'nm']})[0]
            xi = d.c2.points
            yi = d.d2.points
            zi = d.channels[channel_index].values
            vmin = data.channels[channel_index].min()
            vmax = data.channels[channel_index].max()
            X, Y, Z = wt.artists.pcolor_helper(xi, yi, zi)
            ax.pcolor(X, Y, Z, vmin=vmin, vmax=vmax, cmap=cm)
            ax.set_xlim(xi.min(), xi.max())
            ax.set_ylim(yi.min(), yi.max())
            wt.artists.corner_text(int(setpoint), fontsize=18, background_alpha=0.33)
            ax.grid()
    # intensity
    subplot_spec = gs[0, 0]
    cax = plt.subplot(gs[0, 1])
    plot(subplot_spec, cax, 0, 'default', yticklabels=True)
    wt.artists.plot_colorbar(cax=cax, cmap='default', label='intensity')
    # color
    subplot_spec = gs[0, 3]
    cax = plt.subplot(gs[0, 4])
    plot(subplot_spec, cax, 1, 'rainbow')
    ticks = np.linspace(1140, 1620, 9)
    wt.artists.plot_colorbar(cax=cax, cmap='rainbow', label='color (nm)', ticks=ticks)
    # y label
    fig.text(0.01, 0.55, '$\mathsf{\Delta}$D2 (deg)', fontsize=18, rotation=90)
    # finish
    plt.savefig('poweramp.png', dpi=300, transparent=True, pad_inches=1)
    plt.close(fig)