#!/usr/bin/python

import numpy as np
import cPickle as pickle
import matplotlib.pyplot as plt

def points():
    data_dict = {}

    f = open('./datasets/flange_R_Z','r')
    (flanger,flangez) = pickle.load(f)
    f.close()
    data_dict['flange'] = [flanger,flangez]

    f = open('./datasets/PA_in_R_Z','r')    
    (PA_in_R,PA_in_Z) = pickle.load(f)
    f.close()
    data_dict['PA_rib'] = [PA_in_R,PA_in_Z]

    f = open('./datasets/PA_in_sensor_locations','r')
    (PA_s_R,PA_s_Z) = pickle.load(f)
    f.close()
    data_dict['PA_sensors'] = [PA_s_R,PA_s_Z]    

    f = open('./datasets/shell_R_Z','r')
    (shellr,shellz) = pickle.load(f)
    f.close()
    data_dict['shell'] = [shellr,shellz]

    f = open('./datasets/chamber_R_Z','r')
    (chamberr,chamberz) = pickle.load(f)
    f.close()
    data_dict['chamber'] = [chamberr,chamberz]

    f = open('./datasets/limiter_R_Z','r')
    ([new_limiterr,new_limiterz],
     [top_limiterr,top_limiterz],
     [bot_limiterr,bot_limiterz],
     [out_limiterr,out_limiterz]) = pickle.load(f)
    f.close()
    data_dict['limiter'] = [new_limiterr,new_limiterz]

    return(data_dict)

def flux_plot_background(fig = None,ax = None,chamber = False,flange = False):
    if fig == None:
        fig = plt.figure()
    if ax == None:
        ax = fig.add_subplot(111,aspect = 'equal')
    bckgd = points()

    clr = {'PA_rib':'k','PA_sensors':'k','flange':'r','chamber':'k','shell':'k','limiter':'k'}
    lw = {'PA_rib':1,'PA_sensors':1,'flange':3,'chamber':2,'shell':2,'limiter':1}
    #fig = plt.figure()
    
    for key,value in bckgd.items():
        if key == 'chamber' and not chamber:
            continue
        if key == 'flange' and not flange:
            continue
        plt.plot(value[0],value[1],clr[key],linewidth = lw[key])
