"""OPA interferometry"""


# --- import --------------------------------------------------------------------------------------


import numpy as np

import WrightTools as wt 

import matplotlib
import matplotlib.pyplot as plt

import os


# --- define --------------------------------------------------------------------------------------


folderpath = os.getcwd()
datafolder = os.path.join(folderpath, '2014.04.10 data')


# --- get data ------------------------------------------------------------------------------------


p = 'OPA_interferometry.wt5'
if False:
    col = wt.Collection(name='root')
    # autointerference
    filepath = os.path.join(datafolder, 'auto interferometry (1) 0 ms wait [1x201] Delay.dat')
    ignore=['w3', 'dref', 'm0', 'm1', 'm2', 'm3', 'm4', 'm5', 'm6']
    d = wt.data.from_COLORS(filepath, name='auto', ignore=ignore, parent=col)    
    # crossinterference
    filepath = os.path.join(datafolder, 'cross interferometry (2) 0 ms wait [1x201] Delay.dat')
    ignore=['w3', 'dref', 'm0', 'm1', 'm2', 'm3', 'm4', 'm5', 'm6']
    d = wt.data.from_COLORS(filepath, name='cross', ignore=ignore, parent=col)
    # interference over time
    filepath = os.path.join(datafolder, 'high res (1) 1 shots 100 ms wait [1000] Capture.dat')
    arr = np.genfromtxt(filepath).T
    xi = np.arange(1000)*.1 # convert to seconds
    yi = arr[8][:256]
    zi = arr[21]
    zi.shape = (-1, 256)
    d = wt.Data(name='time', parent=col)
    d.create_variable('time', values=xi[:, None])
    d.create_variable('wa', values=yi[None, :], units='nm', label='a')
    d.create_channel('array', values=zi)
    d.transform('time', 'wa')
    col.save(p, overwrite=True)
   
col = wt.open(p)
    
    
# --- auto and cross correlation figure -----------------------------------------------------------


if True:
    def plotter(ax, d, label):
       ax.pcolor(d, channel='mc') # pcolor is kinda distracting
       ax.set_ylim(1200,1400)
       ax.set_xlabel('delay (fs)')
       ax.grid()
       wt.artists.corner_text(label)
       
    
    fig, gs = wt.artists.create_figure(width='double', cols=[1,1,'cbar'])
    # auto
    ax = plt.subplot(gs[0,0])
    plotter(ax, col.auto, 'auto-interference')
    # cross
    ax = plt.subplot(gs[0,1])
    plotter(ax, col.cross, 'cross-interference')
    # pretty up
    wt.artists.set_fig_labels(ylabel=col.auto.wa.label)
    # colorbar
    cax = plt.subplot(gs[0,-1])
    ticks = np.linspace(0,1,11)
    label = 'norm array signal'
    wt.artists.plot_colorbar(cax, ticks=ticks, label=label)
    # finish
    p = 'auto_cross_interference.png'
    wt.artists.savefig(p)


# --- over time -----------------------------------------------------------------------------------


if True:
    def plotter(ax, d):
       ax.pcolor(d) # pcolor is kinda distracting
       ax.set_ylim(1250,1350)
       ax.set_xlabel('lab time (s)')
       ax.set_ylabel(col.auto.wa.label)
       ax.grid()
    d.array.clip(.3)
    fig, gs = wt.artists.create_figure(width='double', cols=[1, 'cbar'], default_aspect=.5)
    # overtime
    ax = plt.subplot(gs[0,0])
    plotter(ax, col.time)
    ax.hlines(1300, 0, 99.9, color='grey', linewidth=5, alpha=.7, zorder=10)
    # colorbar
    cax = plt.subplot(gs[0,-1])
    ticks = np.linspace(0,1,11)
    label = 'norm array signal'
    wt.artists.plot_colorbar(cax, ticks=ticks, label=label)
    # finish
    p = 'time_interference.png'
    wt.artists.savefig(p)