""" OPA power and ranges """

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


import os
import collections

import numpy as np

import matplotlib
import matplotlib.pyplot as plt

import WrightTools as wt
wt.artists.apply_rcparams('publication')


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


p = 'Energy_10742.csv'
idler = np.genfromtxt(p, usecols=(0,1), delimiter=',', skip_header=2, unpack=True)
signal = np.genfromtxt(p, usecols=(2,3), delimiter=',', skip_header=2, unpack=True)
shi = np.genfromtxt(p, usecols=(4,5), delimiter=',', skip_header=2, unpack=True)
shs = np.genfromtxt(p, usecols=(6,7), delimiter=',', skip_header=2, unpack=True)
sfi = np.genfromtxt(p, usecols=(8,9), delimiter=',', skip_header=2, unpack=True)
sfs = np.genfromtxt(p, usecols=(10,11), delimiter=',', skip_header=2, unpack=True)

names = ['idler', 'signal', 'SHI', 'SHS', 'SFI', 'SFS']
wavelengths = [idler[0], signal[0], shi[0], shs[0], sfi[0], sfs[0]]
wavelengths_ev = [wt.units.converter(w, 'nm', 'eV') for w in wavelengths]
energys = [idler[1], signal[1], shi[1], shs[1], sfi[1], sfs[1]]
colors = ['C1', 'C2', 'C3', 'C4', 'C5', 'C6']


# --- workspace -----------------------------------------------------------------------------------


if True:
    # start
    # cbar is shady way to fit legend in
    fig, gs = wt.artists.create_figure(width='single', nrows=1, cols=[1, 'cbar'], default_aspect=1)
    ax = plt.subplot(gs[0,0])
    ax.semilogy( wavelengths_ev[0], energys[0], alpha=0)
    for n,x,y,c in zip(names, wavelengths_ev, energys, colors):
        ax.scatter(x,y, color=c, label=n, marker='o')
    # pretty up    
    ax.grid(which='both')
    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), frameon=False)
    ax.set_xlabel('OPA output (eV)', fontsize=18)
    ax.set_ylabel(u'pulse energy ($\mu$J)')
    # finish
    wt.artists.savefig('OPA_powers.png')

        
# --- old workspace -------------------------------------------------------------------------------


if True:
    # prepare figure
    fig, gs = wt.artists.create_figure(width='single', nrows=1, cols=[1], default_aspect=0.3)
    
    # ranges
    ax = plt.subplot(gs[0, 0])
    cs = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9']
    ranges = collections.OrderedDict()
    ranges['DFG'] = ([11000, 2600], 1)
    ranges['Idl'] = ([2600, 1600], 0.5)
    ranges['Sig'] = ([1600, 1150], 0)
    ranges['SHI'] = ([800, 1200], 0.5)
    ranges['SHS'] = ([800, 580], 0)
    ranges['SFI'] = ([600, 533], 0.5)
    ranges['SFS'] = ([540, 480], 0)
    ranges['4HI'] = ([590, 400], 1)
    ranges['4HS'] = ([400, 290], 0)
    height = 0
    for i, item in enumerate(ranges.items()):
        name, value = item
        limits, height = value
        limits = [wt.units.converter(n, 'nm', 'eV') for n in limits]
        ax.plot(limits, [height, height], lw=25, label=name, c=cs[i], solid_capstyle='butt')
        ax.text(np.mean(limits), height, name, va='center', ha='center', fontsize=11)
    
    ranges = collections.OrderedDict()
    ranges['UVB'] = ([315, 280], wt.artists._colors.nm_to_rgb(400))
    ranges['UVA'] = ([400, 315], wt.artists._colors.nm_to_rgb(450))
    ranges['visible'] = ([700, 400], wt.artists._colors.nm_to_rgb(500))
    ranges['NIR'] = ([2500, 700], wt.artists._colors.nm_to_rgb(600))
    ranges['IR'] = ([15000, 2500], wt.artists._colors.nm_to_rgb(700))
    for i, item in enumerate(ranges.items()):
        name, value = item
        limits, c = value
        limits = [wt.units.converter(n, 'nm', 'eV') for n in limits]
        ax.axvline(limits[0], c='k', lw=1, zorder=1)
        ax.axvline(limits[1], c='k', lw=1, zorder=1)
        ax.text(np.mean(limits), 1.25, name, va='bottom', ha='center', fontsize=12)
        for limit in limits:
            nm = wt.units.converter(limit, 'eV', 'nm')
            ax.text(limit, 1.6, '%i'%nm, fontsize=14, va='bottom', ha='center', rotation=90)
        
        width = limits[1]-limits[0]
        patch = matplotlib.patches.Rectangle((limits[0], -1), width, 10, facecolor=c, alpha=0.1)
        ax.add_patch(patch)
        
    ax.set_yticks([], [])
    ax.set_xlim(0, 4.5)
    ax.set_ylim(-0.25, 1.5)
    ax.set_xlabel('OPA output (eV)', fontsize=18)
    ax.grid(ls=':')
    ax.set_title('OPA output (nm)', fontsize=18, y=1.3)
    # finish
    wt.artists.savefig('OPA_ranges.png')