text: ## @package createShotData
  # Contains all the routines necessary for converting
  # experimental data in a MDSplus tree into a psi-tet
  # dictionary
  import MDSplus
  from psitet import psiObject
  from map_probes import sp_name_dict, dead_probes
  from scipy.io import savemat
  from dataclasses import asdict
  import numpy as np
  
  ## Reads in experimental shot information
  # from a MDSplus tree and reformats into a
  # psi-tet dictionary. Works for HIT-SI and HIT-SI3
  # @param shotname The name of the experimental discharge (string)
  def createShotData(shotname):
      if(len(shotname) == 6):
          treetype = 'hitsi'
      elif(len(shotname) == 9):
          treetype = 'hitsi3'
      else:
          print('Shotname does not match any known tree')
          exit()
      shot = int(shotname)
      corr_probe_shot = 190604001 # first shot with correct R_CAB in tree
      if shot > 190604001:
          corr_probe = 0 #Don't correct
      elif shot < 181129001:
          corr_probe = 1 # Need to correct for cable resistance
          print('pre-Nov. 29th, 2018 data, cable resistance correction WILL be added');
      else:
          corr_probe = 2 #Need to correct data that was double corrected.
          print('shot between Nov. 29th, 2018 and Jun. 4th, 2019: probe resistance overcorrection WILL be fixed');
  
      p = asdict(psiObject())
  
      t0 = 0
      tf = 4.5e-3
      dt = 1.995001e-6
      tsize = int(tf/dt)+1
      time = np.linspace(t0,tf,tsize)
  
      probe = ['01', '02', '03', '04', '05', '06', \
          '07', '08', '09', '10', '11', '12', \
          '13', '14', '15', '16', '17']
      array = 'M'
  
      tree = MDSplus.Tree(treetype, \
          np.asarray(shotname).astype(int), \
          'READONLY')
      p['time'] = time
      p['freq'] = tree.getNode('\\SIHI_FREQ').data()
      if treetype == 'hitsi':
          p['curr01'] = gen_data_in(tree,'i_inj_x',time)
          p['curr02'] = gen_data_in(tree,'i_inj_y',time)
          p['flux01'] = gen_data_in(tree,'psi_inj_x',time)
          p['flux02'] = gen_data_in(tree,'psi_inj_y',time)
          p['v_divfx01'] = gen_data_in(tree,'v_divfx',time)
          p['v_divfy02'] = gen_data_in(tree,'v_divfy',time)
          p['volt01'] = gen_data_in(tree,'v_inj_x',time)
          p['volt02'] = gen_data_in(tree,'v_inj_y',time)
          p['tcurr'] = gen_data_in(tree,'i_tor_spaavg',time)
          p['is_HITSI3'] = False
      elif treetype == 'hitsi3':
          p['curr01'] = gen_data_in(tree,'i_inj_a',time)
          p['curr02'] = gen_data_in(tree,'i_inj_b',time)
          p['curr03'] = gen_data_in(tree,'i_inj_c',time)
          p['flux01'] = gen_data_in(tree,'psi_inj_a',time)
          p['flux02'] = gen_data_in(tree,'psi_inj_b',time)
          p['flux03'] = gen_data_in(tree,'psi_inj_c',time)
          p['volt01'] = gen_data_in(tree,'v_inj_a',time)
          p['volt02'] = gen_data_in(tree,'v_inj_b',time)
          p['volt03'] = gen_data_in(tree,'v_inj_c',time)
          p['tcurr'] = gen_data_in(tree,'i_tor_spaavg',time)
          p['is_HITSI3'] = True
      else:
          print("not a valid Tree")
  
      try:
          p['inter_n'] = gen_data_in(tree,'n_avg_s1',time)
          pass
      except MDSplus.TreeNODATA:
          print('no FIR signal')
          pass
  
      cb_field, ins_depth = get_imp(tree,time,probe,array)
  
      p['modes_mag'],p['modes_phase'] = get_modes(tree,time)
  
      p['sp_Bpol'],p['sp_Btor'],p['sp_names'],p['sp_B'] = \
          get_sp(tree,time)
  
      p['imp_Brad'],p['imp_Bpol'],p['imp_Btor'] = \
          imp_correction(tree,shot,corr_probe, \
              corr_probe_shot,probe,array,cb_field)
  
      # write everything to file
      filename = 'exppsi_'+shotname+'.mat'
      savemat(filename,p)
  
  ## Gets the IMP signals
  # @param tree A MDSplus tree object
  # @param time The surface probe time base
  # @param probe List of probe names
  # @param array Not sure, but related to
  # which IMP signals to read
  # @returns cb_field Uncorrected IMP magnetic field vector
  # @returns ins_depth Insert depth of the IMP
  def get_imp(tree,time,probe,array):
      tsize = len(time)
      shift = True
      dafi = 'dafi_cf'
      shot = tree.shot
  
      N = len(probe)
  
      cb_field = np.zeros((3, N, tsize))
      ins_depth = np.zeros(N)
  
      # calibration factors based on dafi impedence.
      # these values good from shot 121973
      pol_scale = [1.02,1.016,1.014,1.018,1.013,1.022, \
          1.027,1.024,1.025,1.022,1.02,1.028,1.029,1.022, \
          1.022,1.024,1.021]
      #there is no first toroidal probe
      tor_scale = [0,1.025,1.012,1.014,1.008,1.022, \
          1.024,1.024,1.023,1.012,1.019,1.02,1.022, \
          1.019,1.021,1.026,1.025]
  
      # insert IMP stuff here
      nodeflags = [True, True, True]
      for j in range(N):
          pnode = 'B_IMP_M_P' +probe[j] +':b_winding'
          tnode = 'B_IMP_M_T' +probe[j] +':b_winding'
          rnode = 'B_IMP_M_R' +probe[j]
          Pnode = tree.getNode('\\'+pnode)
          Tnode = tree.getNode('\\'+tnode)
          Rnode = tree.getNode('\\'+rnode)
  
          try:
              rsig = Rnode.data()
              pass
          except MDSplus.TreeNODATA:
              nodeflags[0] = False
              cbw_rad = np.zeros(len(time))*np.nan
              pass
          try:
              psig = Pnode.data()
              pass
          except MDSplus.TreeNODATA:
              nodeflags[1] = False
              cbw_pol = np.zeros(len(time))*np.nan
              pass
          try:
              tsig = Tnode.data()
              pass
          except MDSplus.TreeNODATA:
              nodeflags[2] = False
              cbw_tor = np.zeros(len(time))*np.nan
              pass
  
          if nodeflags[0] == True:
              dtr = tree.tdiExecute('samplinginterval(\\'+rnode+')')
              tminr = tree.tdiExecute('minval(dim_of(\\'+rnode+'))')
              tlengthr = len(tree.tdiExecute('dim_of(\\'+rnode+')').data())
              trad = tminr + dtr*np.linspace(0,tlengthr,tlengthr)
              # shifting time base here for digi differences
              if shift:
                  trad = imp_time_shift(trad, shot, array, \
                      probe[j], 'R')
              bw_rad = tree.tdiExecute( \
                  '''slanted_baseline2(sub_baseline_string("\\\\''' +rnode+r'"))')
  
              cbw_rad = np.interp(time,trad,bw_rad)
  
          if nodeflags[1] == True:
              dtp = tree.tdiExecute('samplinginterval(\\'+pnode+')')
              tminp = tree.tdiExecute('minval(dim_of(\\'+pnode+'))')
              tlengthp = len(tree.tdiExecute('dim_of(\\'+pnode+')').data())
              tpol = tminr + dtp*np.linspace(0,tlengthp,tlengthp)
              if shift:
                  tpol = imp_time_shift(tpol, shot, array, \
                      probe[j], 'P')
              bw_pol = tree.tdiExecute( \
                  '''slanted_baseline2(sub_baseline_string("\\\\''' +pnode+r'"))')
  
              cbw_pol = np.interp(time,tpol,bw_pol)
              if dafi == 'dafi_cf':
                  if shot >= 121973:
                      cbw_pol = cbw_pol*pol_scale[j]
  
          if nodeflags[2] == True:
              dtt = tree.tdiExecute('samplinginterval(\\'+tnode+')')
              tmint = tree.tdiExecute('minval(dim_of(\\'+tnode+'))')
              tlengtht = len(tree.tdiExecute('dim_of(\\'+tnode+')').data())
              ttor = tmint + dtr*np.linspace(0,tlengtht,tlengtht)
              # shifting time base here for digi differences
              if shift:
                  ttor = imp_time_shift(ttor, shot, array, \
                      probe[j], 'T')
              bw_tor = tree.tdiExecute( \
                  '''slanted_baseline2(sub_baseline_string("\\\\''' +tnode+r'"))')
  
              cbw_tor = np.interp(time,ttor,bw_tor)
              if dafi == 'dafi_cf':
                  if shot >= 121973:
                      cbw_tor = cbw_tor*tor_scale[j]
  
          cb_field[0,j,:] = cbw_rad
          if array == 'M':
              if j == 1:
                  # there is no rot ang for the 1st probe
                  # b/c there is no toroidal probe
                  cb_field[1,j,:] = cbw_pol
                  cb_field[2,j,:] = cbw_tor
              elif j > 1:
                  rotnode = 'B_IMP_'+array+'_T'+ \
                      probe[j]+':ROT_ANG'
                  rot_ang = tree.getNode('\\'+rotnode).data()
                  cb_field[1,j,:] = cbw_pol*np.cos(rot_ang) - \
                      cbw_tor*np.sin(rot_ang)
                  cb_field[2,j,:] = cbw_tor*np.cos(rot_ang) + \
                      cbw_pol*np.sin(rot_ang)
          else:
              rotnode = 'B_IMP_'+array+'_T'+probe[j]+':ROT_ANG'
              try:
                  rot_ang = tree.getNode('\\'+rotnode).data()
                  cb_field[1,j,:] = cbw_pol*np.cos(rot_ang) - \
                      cbw_tor*np.sin(rot_ang)
                  cb_field[2,j,:] = cbw_tor*np.cos(rot_ang) + \
                      cbw_pol*np.sin(rot_ang)
                  pass
              except MDSplus.TreeNODATA:
                  cb_field[1,j,:] = cbw_pol
                  cb_field[2,j,:] = cbw_tor
                  pass
          r_string1 = r'\B_IMP_M_R'+probe[j]+ \
              r':R:R_CAL_FACT'
          r_string2 = r'\b_imp_ins_d'
          ins_depth[j] = tree.getNode(r_string1).data() - \
              tree.getNode(r_string2).data()
  
      #imp_Brad = cb_field[0,:,:]
      #imp_Bpol = cb_field[1,:,:]
      #imp_Btor = cb_field[2,:,:]
      # 'r' correction
      if shot >= 150122011 and shot < 151112006:
          ins_depth = ins_depth + 0.076
  
      return cb_field,ins_depth
  
  ## Gets the surface probe signals
  # @param tree A MDSplus tree object
  # @param time The surface probe time base
  # @returns sp_Bpol Poloidal surface probe signals
  # @returns sp_Btor Toroidal surface probe signals
  # @returns sp_names surface probe names
  def get_sp(tree,time):
      jp = 0
      sp_Bpol = []
      sp_Btor = []
      sp_names = []
      sp_B = []
      for node in sp_name_dict.keys():
          if node in dead_probes:
              continue
          sp_names.append(node)
          #node = 'b' + node[1:]
          if node[5] == 'P':
              sp_Bpol.append(gen_data_in(tree,node,time))
          else:
              sp_Btor.append(gen_data_in(tree,node,time))
          sp_B.append(gen_data_in(tree,node,time))
      return sp_Bpol,sp_Btor,sp_names,sp_B
  
  ## Shifts the time base of the probe signals due to time
  # base differences between digitizers
  # @param time The surface probe time base
  # @param shot Shot number (integer)
  # @param array Not sure, but related to
  # which IMP signals to read
  # @param probe List of probe names
  # @param dir magnetic field direction, R, P, T
  # @returns tout The new time base for the IMP
  def imp_time_shift(time,shot,array,probe,dir):
  
      shift612 = 5e-6
      shift2412 = 2.5e-6
      tout = time
  
      if shot < 117860:
          tout = tin
  
      elif shot >= 117860 and shot <= 118389:
          if array == 'M':
              if probe == '06':
                  if dir == 'R':
                      tout = time - shift2412
          elif array == 'B':
              if probe == '02':
                  tout = time - shift612
              elif probe == '03':
                  tout = time - shift612
              elif probe == '04':
                  tout = time - shift2412
              elif probe == '05':
                  tout = time - shift612
              elif probe == '06':
                  tout = time - shift612
  
      elif shot > 118389 and shot <= 121973:
          if array == 'M':
              if probe == '08':
                  if dir == 'R':
                      tout = time - shift2412
              elif probe == '10':
                  tout = time - shift612
              elif probe == '12':
                  tout = time - shift612
              elif probe == '14':
                  tout = time - shift612
              elif probe == '17':
                  tout = time - shift612
          elif array == 'B':
              if probe == '08':
                  tout = time - shift2412
  
      elif shot > 121973 and shot <= 127542:
          if array == 'M':
              if probe == '06':
                  if dir == 'P':
                      tout = time - shift2412
                  elif dir == 'T':
                      tout = time - shift2412
              elif probe == '15':
                  if dir == 'P':
                      tout = time - shift2412
                  elif dir == 'T':
                      tout = time - shift2412
              if dir == 'R':
                  tout = time - shift612
  
      # changed the digitization rate before these shots so the time shift is
      # smaller (only the 612's are off by as much as a usec)
      elif shot > 127542:
          tout = time
  
      return tout
  
  ## Helper function for interpolating onto surface
  # probe time base
  # @param tree A MDSplus tree object
  # @node Name of a valid MDSplus tree node (string)
  # @param time The surface probe time base
  # @returns x The signal corresponding to the node in the Tree,
  # interpolated to the surface probe timebase
  def gen_data_in(tree,node,time):
      x = tree.getNode('\\'+node).data()
      dt = tree.tdiExecute('samplinginterval(\\'+node+')')
      tmin = tree.tdiExecute('minval(dim_of(\\'+node+'))')
      tlength = len(x)
      t = tmin + dt*np.linspace(0,tlength,tlength)
      x = np.interp(time,t,x)
      return x
  
  ## Helper function getting the fourier modes from
  # the diagnostic gap probes
  # @param tree A MDSplus tree object
  # @param time The surface probe time base
  # @returns modes_mag Magnitude of the fourier modes
  # @returns modes_phase Phases of the fourier modes
  def get_modes(tree,time):
      tsize = len(time)
      modes_mag = np.zeros((4,7,tsize))
      modes_phase = np.zeros((4,7,tsize))
      for i in range(7):
          node = 'b_t_l05_n'+str(i)+':magnitude'
          modes_mag[0,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_t_l06_n'+str(i)+':magnitude'
          modes_mag[1,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_p_l05_n'+str(i)+':magnitude'
          modes_mag[2,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_p_l06_n'+str(i)+':magnitude'
          modes_mag[3,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_t_l05_n'+str(i)+':phase'
          modes_phase[0,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_t_l06_n'+str(i)+':phase'
          modes_phase[1,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_p_l05_n'+str(i)+':phase'
          modes_phase[2,i,:] = \
              gen_data_in(tree,node,time)
          node = 'b_p_l06_n'+str(i)+':phase'
          modes_phase[3,i,:] = \
              gen_data_in(tree,node,time)
      return modes_mag,modes_phase
  
  ## This function accounts for IMP corrections beyond
  # the simple DAFI corrections.
  # This is based on imp_in.m from ACH on June 4 2019
  # @param tree A MDSplus tree object
  # @param shot Shot number (integer)
  # @param corr_probe The surface probe time base
  # @param corr_probe_shot Shot name of the hitsi3 reference
  # with the correct IMP values
  # @param probe List of probe names
  # @param array Not sure, but related to
  # which IMP signals to read
  # @param cb_field Uncorrected IMP magnetic field vector
  # @returns imp_Brad Corrected radial IMP signals
  # @returns imp_Bpol Corrected poloidal IMP signals
  # @returns imp_Btor Corrected toroidal IMP signals
  def imp_correction(tree,shot,corr_probe,corr_probe_shot,probe,array,cb_field):
      dir = ['R','P','T']
      N = len(probe)
      if corr_probe==1:
          tree.close()
          new_tree = MDSplus.Tree('hitsi3',corr_probe_shot,'READONLY')
          R_P = np.zeros((3,N))
          for n in range(3):
              for m in range(N):
                  nodeR_P = 'B_IMP_'+array+'_'+dir[n] \
                      +probe[m]+':R_CAB' # probe resistance
                  try:
                      R_P = new_tree.getNode('\\'+nodeR_P).data()
                      pass
                  except MDSplus.TreeNODATA:
                      R_P = 8
                      print('No R_CAB data for node: '+ \
                          nodeR_P+', assuming R_P = 8')
                      pass
  
                  R_T = 50.0 # DAFI variations from 50 Ohms already accounted for
                  calFac = (R_T + R_P) / R_T
                  cb_field[n,m,:] = calFac * cb_field[n,m,:]
      elif corr_probe == 2:
          incorrectR_P = np.zeros((3, N))
          for n in range(3):
              for m in range(N):
                  nodeR_P = 'B_IMP_'+array+'_'+ \
                      dir[n]+probe[m]+':R_P'
                  try:
                      R_P = tree.getNode('\\'+nodeR_P).data()
                      pass
                  except MDSplus.TreeNODATA:
                      R_P = 8
                      print('No R_CAB data for node: '+ \
                          nodeR_P+', assuming R_P = 8')
                      pass
                  incorrectR_P[n,m] = R_P
  
          tree.close()
          new_tree = MDSplus.Tree('hitsi3',corr_probe_shot,'READONLY')
          R_P = np.zeros((3, N))
          for n in range(3):
              for m in range(N):
                  nodeR_CAB = 'B_IMP_'+array+'_'+ \
                      dir[n]+probe[m]+':R_CAB'
                  try:
                      R_CAB = new_tree.getNode('\\'+nodeR_CAB).data()
                      pass
                  except MDSplus.TreeNODATA:
                      R_CAB = 4
                      print('No R_CAB data for node: '+ \
                          nodeR_CAB+', assuming R_P = 4')
                      pass
                  R_P[n,m] = R_CAB
  
          R_T = 50.0
          calFac = (R_P + R_T) / (incorrectR_P + R_T)
          for n in range(np.shape(cb_field)[2]):
              cb_field[:,:,n] = calFac * cb_field[:,:,n]
  
      imp_Brad = cb_field[0,:,:]
      imp_Bpol = cb_field[1,:,:]
      imp_Btor = cb_field[2,:,:]
  
      return imp_Brad, imp_Bpol, imp_Btor
  
  