from pygolem_lite import Shot, saveconst
from pygolem_lite.modules import *
from pygolem_lite.utilities import read_config

import numpy as np

import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
matplotlib.rc('font',  size='10')

import matplotlib.pyplot as plt
import os, sys


R = 47.0                                   # [Ohm] resistance for U -> I


def calculate_M(U_up, U_down):
    return 0.4 * np.log(U_up / U_down) # should be I, but U/R / U/R would be the same


def make_img1(t, M, up_ch, down_ch):
    plt.figure(figsize=[12, 6])          #size in inches
    plt.plot(t * 1e3, M)                  # in ms
    plt.xlabel('htime [ms]')
    plt.ylabel('M')
    fname = 'M1-ch{}_ch{}.png'.format(up_ch, down_ch)
    plt.savefig(fname)
    plt.close()                           #close this figure

def make_img2(t, M, up_ch, down_ch):
    plt.figure(figsize=[12, 6])          #size in inches
    plt.plot(t * 1e3, M)                  # in ms
    plt.xlabel('time [ms]')
    plt.ylabel('M')
    fname = 'M2-ch{}_ch{}.png'.format(up_ch, down_ch)
    plt.savefig(fname)
    plt.close()                           #close this figure


def make_plots1(probe_channels1, automatic_select_up):
    shot = Shot()
    t, data_st = shot['papouch_st']

    start = shot['plasma_start']
    end = shot['plasma_end']
    start_idx = np.where(start <= t)[0][0]
    end_idx = np.where(t <= end)[0][-1]
    t = t[start_idx:end_idx]

    for up_ch, down_ch in probe_channels1:
        U_down = data_st[start_idx:end_idx,down_ch-1] # data columns is indexed from 0
        U_up = data_st[start_idx:end_idx,up_ch-1]
       # if automatic_select_up and U_down.mean() > U_up.mean():
          #  U_up, U_down = U_down, U_up   #swap
        M = calculate_M(U_up, U_down)
        make_img1(t, M, up_ch, down_ch)

        
probe_channels1 = [
    (1, 2),
    (3, 4),
    (5, 6),
    (7, 8),
    (9,10),
    (11,12),
    ]

def make_plots2(probe_channels2, automatic_select_up):
    shot = Shot()
    t, data_ja = shot['papouch_ja']

    start = shot['plasma_start']
    end = shot['plasma_end']
    start_idx = np.where(start <= t)[0][0]
    end_idx = np.where(t <= end)[0][-1]
    t = t[start_idx:end_idx]

    for up_ch, down_ch in probe_channels2:
        U_down = data_ja[start_idx:end_idx,down_ch-1] # data columns is indexed from 0
        U_up = data_ja[start_idx:end_idx,up_ch-1]
       # if automatic_select_up and U_down.mean() > U_up.mean():
          #  U_up, U_down = U_down, U_up   #swap
        M = calculate_M(U_up, U_down)
        make_img2(t, M, up_ch, down_ch)

        
probe_channels2 = [
    (1, 2),
    (3, 4),
    ]
    
def graphs(file_type):
    shot = Shot()
    ShotNo=shot['shotno']
    os.system("cp ../../../DAS/1011Papouch_St.ON/*.npz .")
    os.system("cp ../../../DAS/0314Papouch_Ja.ON/*.npz .")

    outUp = [get_data('Papouch_St_01.npz', "M1" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_03.npz', "M2" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_05.npz', "M3" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_07.npz', "M4" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_Ja_04.npz', "M5" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_09.npz', "M6" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_11.npz', "M7" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_Ja_01.npz', "M8" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False) ]
    
    
#	multiplot(out, name  , 'graph', (6,len(out)*2), 100, 'vertical', file_type)

    multiplot(outUp, '\#'+str(int(ShotNo))+' Up'  , 'graphUp', (6,len(outUp)*2), 100, 'vertical', file_type)

    outDown = [get_data('Papouch_St_02.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_04.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_06.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_08.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_Ja_03.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_10.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_St_12.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False),get_data('Papouch_Ja_02.npz', "" , "",xlim=[9,18],ylim=[-0.5,0.5] ,integrated=False, reduction = False) ]
    
    multiplot(outDown, '\#'+str(int(ShotNo))+' Down'  , 'graphDown', (6,len(outDown)*2), 100, 'vertical', file_type)
	
if __name__ == '__main__':
    make_plots1(probe_channels1, automatic_select_up=True)
    make_plots2(probe_channels2, automatic_select_up=True)
#    if sys.argv[1] ==  "plots":
    graphs('png')
    saveconst('status', 0)