import re
import numpy as np
import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
import matplotlib.pyplot as plt
from pygolem_lite import Shot

def get_r_a1(shot):
    comment = shot['wwwcomment']
    match = re.match('.*r_a1 = (\\S+) mm.*', comment)
    try:
        return float(match.group(1))
    except Exception as e:
        print e
        return 0


if __name__ == '__main__':
    shot = Shot()
    t, data = shot['papouch_ko']
    try:
        _, data_2 = shot['papouch_za']
        data = np.hstack((data, data_2[:,:4]))
    except Exception as e:
        print e
    t_start = shot['plasma_start']
    t_end = shot['plasma_end']
    stationary_window = np.logical_and(t_start + 5e-3 <= t,
                                       t <= t_end - 3e-3)
    data = data[stationary_window,:]
    channels = data.shape[1]
    a_idx = np.arange(0, channels, 2)
    b_idx = a_idx + 1
    E_pol = data[:,a_idx] - data[:,b_idx]
    E_pol = (E_pol[:,1:] + E_pol[:,:-1]) / 2.0 # center out
    E_rad = (data[:,a_idx] + data[:,b_idx]) / 2.0
    E_rad = E_rad[:,:-1] - E_rad[:,1:]
    R_e = E_rad * E_pol
    R_e_m = np.mean(R_e, axis=0)
    R_e_std = np.std(R_e, axis=0)

    r_a1 = get_r_a1(shot)
    r = r_a1 + np.arange(R_e_m.shape[0]) * 2.5 + 1.25 # centered
    plt.errorbar(r, R_e_m, R_e_std, fmt='--')
    plt.xlabel('$r$ [mm]')
    plt.ylabel('$\\propto \langle E_\\theta E_r \\rangle$')
    plt.grid(True)
    plt.axvline(85, color='r', linestyle='--', linewidth=2, label='limiter')
    plt.savefig('plot.png')
    print np.min(R_e_m)