# SPEC reader
# 2016-11-23
# Sam Tardif (samuel.tardif@gmail.com)
# Python3 compatibility by Nils Blanc

import numpy as np
import os.path
import string


class SpecFile:
  """

  Simple class to read SPEC file. The main purposes are :
  (i) read the header and initiate the relevant properties
  (ii) build a dictionary of the scan number and their binary position in the file
  The latter allows for faster subsequent scan reading

  Definition:
  -----------
  SpecFile(spec_file, verbose = False)
   > spec_file : string
   > scan_numbers : integer or as tuple/list/array of integer

  Attributes (typical):
  -----------
  comments........all comments
  date............scan start datestamp
  file............string of the SPEC file name
  scan_dict.......dictionary {scan number : binary position in file}

  Examples:
  --------
  # read the specfile
  In : sf = SpecFile('./lineup0.dat')

  """
  # dictionary definitions for handling the spec identifiers
  def __param__(self):
    return  { 'S' : self.__scanline__,
              'D' : self.__dating__,
              'T' : self.__counting__,
              'G' : self.__configurating__,
              'Q' : self.__hkl__,
              'O' : self.__motorslabeling__,
              'o' : self.__motorslabelingnospace__,
              'P' : self.__positioning__,
              'M' : self.__marccdpath__,
              'N' : self.__speccol__,
              'L' : self.__counterslabeling__,
              'J' : self.__allcounterslabeling__,
              'j' : self.__allcounterslabelingnospace__,
              'C' : self.__commenting__,
              'E' : self.__initialepoch__,
              'F' : self.__specfilenaming__,
              '@' : self.__special__,
              'XIAFILE' : self.__xiafilenaming__,
              'XIACALIB' : self.__xiacalibrating__,
              'XIAROI' : self.__xiaroidefining__,
              'ULIMA_mpx4' : self.__limampx4path__,
              'UDETCALIB' : self.__detcalib__}

  def __scanline__(self, l):
    self.number = int(l.split()[1])
    self.type = l.split()[2]
    self.args = l.split()[3:]
    self.command = l[3+len(l.split()[1])+2:].rstrip()

  def __dating__(self, l):
    self.date = " ".join(l.split()[1:])

  def __counting__(self, l):
    self.ct = float(l.split()[1])
    self.ct_units = l.split()[2]

  def __configurating__(self, l):
    self.__config__ = self.__config__ + " ".join(l.split()[1:]) + " "

  def __hkl__(self, l):
    self.Qo = l.split()[1:]

  def __motorslabeling__(self, l):
    self.__motorslabels__ = self.__motorslabels__ + " ".join(l.split()[1:]) + " "

  def __motorslabelingnospace__(self, l):
    self.__motorslabelsnospace__ = self.__motorslabelsnospace__ + " ".join(l.split()[1:]) + " "

  def __positioning__(self, l):
    self.__positions__ = self.__positions__ + " ".join(l.split()[1:]) + " "

  def __speccol__(self, l):
    self.N = int(" ".join(l.split()[1:]))

  def __counterslabeling__(self, l):
    # column labels
    self.counters = l.split()[1:]

  def __allcounterslabeling__(self, l):
    # all counters in the experiment
    self.__counters__ = self.__counters__  + " ".join(l.split()[1:]) + " "
    
  def __allcounterslabelingnospace__(self, l):
    # all counters in the experiment
    self.__countersnospace__ = self.__countersnospace__  + " ".join(l.split()[1:]) + " "
  
  def __specfilenaming__(self, l):
    self.specfilename = l.split()[1:]

  def __commenting__(self, l):
    self.comments = self.comments + " ".join(l.split()[1:])  

  def __initialepoch__(self, l):
    self.Epoch0 = int(l.split()[1])

  def __marccdpath__(self, l):
    self.M = l.split()[1:]
    
  def __limampx4path__(self, l):
    self.limampx4path = os.path.join(*l.split()[1:][0].split(os.path.sep))  #l.split()[1:][0]

  def __detcalib__(self, l):
    detcalib = l.split(' ')[1].split(',')
#    cen_pix_x=352.585,cen_pix_y=139.262,pixperdeg=315.152,det_distance_CC=0.993,det_distance_COM=0.992,timestamp=2017-11-10T11:54:52.621448
    self.detcalib_cen_pix_x = np.float(detcalib[0].split('=')[1])
    self.detcalib_cen_pix_y = np.float(detcalib[1].split('=')[1])
    self.detcalib_pixperdeg = np.float(detcalib[2].split('=')[1])
    self.detcalib_det_distance_CC = np.float(detcalib[3].split('=')[1])
    self.detcalib_det_distance_COM = np.float(detcalib[4].split('=')[1])
    self.detcalib_timestamp = detcalib[5].split('=')[1]  
    
  def __special__(self, l):
    self.__param__()[l[2:].split(' ')[0]](l) 
    
  def __xiafilenaming__(self, l):
    self.xianame = l.split()[1:][0]
    self.xiaroi = dict()
    
  def __xiacalibrating__(self, l):
    self.xiacalib = l.split()[1:][0]
    
  def __xiaroidefining__(self, l):
    self.xiaroi[l.split()[1]] = [int(l.split()[2]),int(l.split()[3]),int(l.split()[4]),int(l.split()[5]),int(l.split()[6])]


  def __readSpecLine__(self,l, verbose=False):
    try: 
      SpecKey = l.split()[0][1:]
      if SpecKey[0] != "U": SpecKey = SpecKey.rstrip(string.digits) # remove trailing digits for SpecKeys other than those starting with "U" (User defined?)
      self.__param__()[SpecKey](l)          
    except KeyError:
      if verbose : print("unprocessed line (SpecKey {}): ".format(SpecKey) + l)


  def __init__(self, spec_file, verbose = False):
    # init a bunch of stuff that will also be used by the children class Scan
    self.file = spec_file
    self.__motorslabels__ = "" # list of all motors in the experiment
    self.__motorslabelsnospace__ = "" # list of all motors in the experiment
    self.__positions__ = "" # list the values of all motors
    self.__counters__ = "" # list of all counters in the experiment
    self.__countersnospace__ = "" # list of all counters in the experiment
    self.__config__ = "" # list the values of the UB matrix config
    self.comments = ""
    self.scan_dict={}  # dictionary to store the position in the file of the scans
    try:
      with open(spec_file,'rU') as f:  # the U mode indicates universal line break, essential for accurate counting
        # first read the file header (mostly comments and motors definition)
        # up to the first scan (identified by a line starting with "#S"
        reading_header = True
        while reading_header:
          position_in_file = f.tell() # get the position AHEAD of the scan first line
          l = f.readline()
          if len(l) > 1: # not an empty line
            if l[:2] == '#S':
              reading_header = False
              scan_number = int(l.split()[1])
              if verbose:
                  print("after reading the header, found scan {} at location {}".format(scan_number,position_in_file))
              self.scan_dict[scan_number] = position_in_file
            else:
              self.__readSpecLine__(l, verbose=verbose)
        reading_file = True  # change this switch when the end of file is reached (i.e. read an empty string)
        while reading_file:
          position_in_file = f.tell() # get the position AHEAD of the scan first line
          l = f.readline()
          if len(l) > 1: # not an empty line
            if l[:2] == '#S':
              scan_number = int(l.split()[1])
              if verbose:
                  print("found scan {} at location {}".format(scan_number,position_in_file))
              self.scan_dict[scan_number] = position_in_file
          if l == "":
            reading_file = False
    except IOError:
      print("could not find the file {}".format(spec_file))



class Scan(SpecFile):
  """
  Simple class to read extract scans from SPEC files. All the parameters of the scan and the data are read
  and stored as attributes.
  
  Definition:
  -----------
  Scan(spec_file, scan_numbers, verbose = False)
   > spec_file as SpecFile object or string (in this case a SpecFile object will be instanced)
   > scan_numbers as integer or as tuple/list/array of integer
  
  Attributes:
  -----------
  <countername>...data in counter <countername> (see counters for description)
  number..........scan number of the first scan in the list
  scan_numbers....list of all the scans included
  type............scan type
  args............scan arguments (motor start stop npoints counting)
  date............scan start datestamp
  ct..............scan counting time per point
  Qo..............H K L position at start of scan
  M...............MarCCD image file path
  N...............number of counters
  counters........list of counters
  motors..........dictionary of motors and their initial position
  comments........all comments
  tstart, tend....starting and finishing time
  duration........duration in s
  time_per_point..duration per point

  Examples:
  --------
  # read the specfile
  In : sf = SpecFile('./lineup0.dat')

  # read the scan
  In : scan = Scan(sf, 265)
  
  # read a series of scans
  In : scan = Scan(sf, (265,266,270))
  In : scan = Scan(sf, arange(265,270))
    
  # learn about the motors
  In : scan.motors
  Out:  
  {'Chi': -144.8913,
   'Phi': 176.9778,
   ...
   'xpr3z': -8.8}
  
  # extract the info on a particular motor
  In : scan.motors['tth']
  Out: 35.042
  
  # plot two counters vs each others
  In : plot(scan.th,scan.det/scan.IC1)
  
  
  
  """
  

  
  def __init__(self, spec_file, scan_numbers, verbose = False):
    if type(spec_file) == str:
      spec_file = SpecFile(spec_file)
    self.file = spec_file.file
    try :
      len_scan_number = len(scan_numbers)  # it is a list of scan
    except TypeError:
      scan_numbers = [scan_numbers,]  # it is a simple scan, we make a len 0 list
    scan_number = scan_numbers[0] #for all intents and purposes
    # recover the names of the motors from the header of the SPEC file (i.e. the SpecFile instance)
    self.__motorslabels__ = spec_file.__motorslabels__ # list of all motors in the experiment
    self.__motorslabelsnospace__ = spec_file.__motorslabelsnospace__ # list of all motors in the experiment
    # prepare the scan-specific attributes
    self.__positions__ = "" # list the values of all motors
    self.__config__ = "" # list the values of the UB matrix config
    self.__counters__ = "" # list of all counters in the experiment
    self.__countersnospace__ = "" # list of all counters in the experiment
    self.scan_numbers = scan_numbers
    self.comments = ""


    with open(self.file,'rU') as f:
      # read the first (and possibly only) scan in the list
      # now try to find the scan
      f.seek(spec_file.scan_dict[scan_number])
      l = f.readline()
      if verbose  : print("reading scan " + l)

      # read the scan header
      while l[0] == '#':
        self.__readSpecLine__(l, verbose=verbose)
        l = f.readline()

      # finally read the data (comments at the end are also read and added to the comment attribute)
      data = [list(map(float,l.split()))]
      l = f.readline()
      while l != '\n' and l != '':
        if l[0] == '#' and l[1] != 'C':
          break
        if l[0] == '#': 
            self.__readSpecLine__(l, verbose=verbose)
        else : 
            data.append(list(map(float,l.split())))
        l = f.readline()



      # now get the data for each scan in the list
      if len(scan_numbers) > 0:
        for scan_number in scan_numbers[1:]:
          # now try to find the scan
          f.seek(spec_file.scan_dict[scan_number])
          l = f.readline()
          if verbose  : print("reading scan " + l)

          # check that we actually concatenate similar scans !
          similar_scan  = (l.split()[2] == self.type)
          if len(self.args) > 1 :
            try:
              similar_scan = similar_scan * (l.split()[3]==self.args[0])
            except ValueError:
              similar_scan = False
          if len(self.args) > 5 :
            try:
              similar_scan = similar_scan * (l.split()[6]==self.args[3])
            except ValueError:
              similar_scan = False


          if similar_scan:
            # read pass the scan header
            while l[0] == '#':
              l = f.readline()

            # finally read the data (comments at the end are also read and added to the comment attribute)
            data.append(list(map(float,l.split())))
            l = f.readline()
            while l != '\n' and l != '':
              if l[0] == '#' and l[1] != 'C':
                break
              if l[0] == '#':
                self.__readSpecLine__(l, verbose=verbose)
              else : 
                data.append(list(map(float,l.split())))
              l = f.readline()
          else :
            print("not all scans are the same type")

      # set the data as attributes with the counter name
      for i in range(len(self.counters)):
        #print(i,range(len(self.counters)),self, self.counters[i], np.asarray(data).shape, np.asarray(data), data)
        setattr(self, self.counters[i], np.asarray(data)[:,i])


      # make the motors/positions dictionary
      # usual case
      if len(self.__motorslabels__.split()) == len(self.__positions__.split()) :
        self.motors = dict(zip(self.__motorslabels__.split(), list(map(float,self.__positions__.split()))))
      # when some motors names have spaces and there is a second line (small o) to describe them
      elif len(self.__motorslabelsnospace__.split()) == len(self.__positions__.split()) :
        self.motors = dict(zip(self.__motorslabelsnospace__.split(), list(map(float,self.__positions__.split()))))


      #TEST : attribute-like dictionary
      # removed due to conflicts when a motor was also a counter
      #    for motor in self.motors:
      #                setattr(self, motor, self.motors[motor])


      # small sanity check, sometimes N is diffrent from the actual number of columns
      # which is known to trouble GUIs like Newplot and PyMCA
      if self.N != len(self.counters):
        print("Watch out! There are %i counters in the scan but SPEC knows only N = %i !!"%(len(self.counters),self.N))

      if hasattr(self,'Epoch'):
        self.tstart = self.Epoch[0]
        self.tend   = self.Epoch[-1]
        self.duration = self.tend - self.tstart
        self.time_per_point = self.duration/len(self.Epoch)
#class Scan2D(Scan):
  #def __init__(self, spec_file, scan_number_list, verbose = verbose):
    
