text: # -*- coding: utf-8 -*-
  
  """
  Module pyelegantsdds.sdds 
  =================================================================
  
  A module containing the SDDSCommand class and  SDDS class.
  
  The SDDSCommand class is responsible for creating, editing, etc ... of SDDS
  commands.
  
  The SDDS class uses the SDDSCommand class to build template and composed commands.
  
  """
  
  import os
  import shlex
  import subprocess as subp
  from io import StringIO
  from itertools import count
  from typing import Dict, List
  
  import numpy as np
  import pandas as pd
  from dask import dataframe as dd
  
  from .tools.sddsutils import sddsconvert2ascii, sddsconvert2binary
  
  
  # ================================================================================
  # SDDSCommand class:
  #
  #
  # ================================================================================
  class SDDSCommand:
      """SDDSCommand manupilation class"""
  
      _COMMANDLIST = [
          "plaindata2sdds",
          "sdds2stream",
          "sddsconvert",
          "sddsquery",
          "sddsprocess",
          "sddsplot",
          "sddsprintout",
          "sddsmakedataset",
          "sddsnaff",
          "sddsresdiag",
      ]
  
      def __init__(self, sif):
          self.sif = sif
          self.command = {}
  
      def createCommand(self, command: str, note: str = "", **params) -> None:
          """
          Method to add a command to the command file.
          Generates a dict to reconstruct command string,
          that can be added to the command file.
  
          Parameters:
          ----------
          command     : str
          note        : str
  
          Key Parameters:
          ---------------
          params : dict
  
          Returns:
          --------
          None
          """
          # check if it is a valid Elegant command
          if self.checkCommand(command) == False:
              print("The command {} is not recognized.".format(command))
              return
  
          # init command dict
          thiscom = {}
          thiscom["NAME"] = command
  
          # add command parameters to the command dict
          for k, v in params.items():
              thiscom[k] = v
  
          # add the command dict to the command list
          self.command = thiscom
  
      def checkCommand(self, typename: str) -> bool:
          """
          Check if a command is a valid
          Elegant command.
  
          Parameters:
          -----------
          typename    : str
  
          Returns:
          --------
          boolean
  
          """
          for tn in self._COMMANDLIST:
              if tn == typename.lower():
                  return True
          return False
  
      def clearCommand(self):
          """
          Clear the command list.
          """
          self.command = {}
  
      def getCommand(self, command, **params):
          """Method to get the command.
  
          Parameters
          ----------
          command : dict
  
          Returns
          -------
          str
          command string
          """
          # create command
          self.createCommand(command, **params)
  
          # init command string
          cmdstr = "{} {}".format(self.sif, self.command["NAME"])
  
          # build the command string
          for k, v in params.items():
              if "file" in k.lower():
                  cmdstr += " {}".format(v)
              elif "separator" in k.lower():
                  cmdstr += ' "-separator= {}"'.format(v)
              else:
                  if v is not None:
                      cmdstr += " -{}={}".format(k.split("_")[0], v)
                  else:
                      cmdstr += " -{}".format(k.split("_")[0])
  
          return cmdstr
  
      def runCommand(self, commandstring=None):
          """Run the command.
  
          Parameters
          ----------
          commandstring : str
          """
          if commandstring is None:
              name = self.command.pop("NAME")
              commandstring = self.getCommand(name, **self.command)
          print("Running command {}".format(commandstring))
          subp.run(commandstring, check=True, shell=True)
  
      def get_particles_plain_2_SDDS_command(self, **params) -> str:
          """
          Returns sdds command to turn plain data table, separated
          by a *space* into SDDS particle initial coordinates SDDS file.
          Easy to generate plain data with pandas.
  
          """
          return self.getCommand(
              "plaindata2sdds",
              file_1=params.get("file_1", "temp_plain_particles.dat"),
              file_2=params.get("file_2", "temp_particles_input.txt"),
              inputMode=params.get("inputMode", "ascii"),
              outputMode=params.get("outputMode", "ascii"),
              separator=" ",
              column_1="x,double,units=m",
              column_2="xp,double",
              column_3="y,double,units=m",
              column_4="yp,double",
              column_5="t,double,units=s",
              column_6='p,double,units="m$be$nc"',
              columns_7="particleID,long",
              noRowCount=None,
          )
  
  
  # ================================================================================
  # SDDS CLASS
  # ================================================================================
  
  
  class SDDS:
      """
      Class for interacting with SDDS files.
      """
  
      def __init__(self, sif: str, filename: str, filetype: int):
          assert filetype in [0, 1]
          self.sif = sif
          self._filetype = filetype
          self._filename = filename
          self.columnlist = None
          self.parameterlist = None
          self.commandlist = []
          self.command_history = {}
          self._count_iter = count(start=0, step=1)
  
      @property
      def filename(self):
          return self._filename
  
      @filename.setter
      def filename(self, value):
          self._filename = value
  
      @property
      def filetype(self):
          return self._filetype
  
      @filetype.setter
      def filetype(self, value):
          assert value in [0, 1]
          self._filetype = value
  
      def addCommand(self, command: str, **params) -> None:
          """Add a command to the command list.
  
          Parameters
          ----------
          command : str
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          valid sdds command
          """
          sddscommand = SDDSCommand(self.sif)
          cmdstr = sddscommand.getCommand(command, **params)
  
          # add command to current commandlist
          self.commandlist.append(cmdstr)
  
      def clearCommandList(self, save: bool = True) -> None:
          """
          Clear the command list.
  
          Parameters:
          -----------
          bool: save
          """
          if save and len(self.commandlist) > 0:
              self._addHistory(self.commandlist)
  
          self.commandlist = []
  
      def clearHistory(self):
          """
          Clear command history.
          """
          self.command_history = {}
          self._count_iter = count(start=0, step=1)
  
      def _addHistory(self, cmdlist):
          """
          Add commands to history.
  
          Parameters:
          -----------
          list | str : cmdlist
          """
          _key = next(self._count_iter)
          if not isinstance(cmdlist, list):
              cmdlist = [cmdlist]
          _value = cmdlist
          self.command_history[_key] = _value
  
      def runCommand(self):
          """
          Run all the sddscommands in the command list.
          """
          # print info
          if len(self.commandlist) > 0:
              # add commands to history
              self._addHistory(self.commandlist)
  
              for cmd in self.commandlist:
                  # run command
                  print("Executing : \n{}".format(cmd))
  
                  p = subp.Popen(cmd, stdout=subp.PIPE, shell=True)
                  (output, err) = p.communicate()
                  p_status = p.wait()
  
          else:
              print("No commands entered - nothing to do!")
  
          # clear the commandlist
          self.clearCommandList()
  
      def printHistory(self):
          """Print command history."""
          if len(self.command_history) > 0:
              for i, l in self.command_history.items():
                  print("History key: {}\n".format(i))
                  print("---------------\n")
                  for c in l:
                      print("{}\n".format(c))
  
                  # print extra newline
                  print("\n")
          else:
              print("History is empty.")
  
      def reload_from_history(self, history_idx: int) -> None:
          """Reloads the commandlist from history with index
          history_idx.
  
          Parameters
          ----------
          history_idx : int
                                                                          index of history command to reload
                                                                          use `printHistory()` to get an overview
          """
          self.clearCommandList(save=True)
          self.commandlist = self.command_history.get(history_idx)
  
      def rerun_from_history(self, history_idx):
          """Rerun a history entry.count()
  
          Parameters
          ----------
          history_idx : int
          """
          assert history_idx in list(self.command_history.keys())
  
          self.clearCommandList(save=True)
          self.commandlist = self.command_history.get(history_idx)
          self.runCommand()
  
      def load_raw_data(self):
          """Read the data from file."""
          if self.filetype == 1:
              # ASCII FORMAT
              with open(self.filename, "r") as f:
                  self.raw_content = f.read()
  
          else:
              # BINARY FORMAT
              with open(self.filename, "rb") as f:
                  self.raw_content = f.read()
  
      def convert(self, outfile: str = None):
          """Convert between filetypes.
  
          Parameters
          ----------
          outfile : str
          """
          if self.filetype == 0:
              converted_filename = self.filename + ".txt"
  
              if outfile is not None:
                  converted_filename = outfile
  
              cmdstr = "{} sddsconvert -ascii {} {}".format(
                  self.sif, self.filename, converted_filename
              )
          else:
              converted_filename = self.filename + ".bin"
  
              if outfile is not None:
                  converted_filename = outfile
  
              cmdstr = "{} sddsconvert -binary {} {}".format(
                  self.sif, self.filename, converted_filename
              )
          # add to command history
          self._addHistory(cmdstr)
  
          # reset filename
          print("Warning - auto filename set")
          print("Changed from {} to {}".format(self.filename, converted_filename))
          self.filename = converted_filename
  
          # reset filetype
          print("Warning - auto filetype set")
          print("Changed from {} to {}".format(self.filetype, abs(1 - self.filetype)))
          self.filetype = abs(1 - self.filetype)
  
          with open(os.devnull, "w") as f:
              subp.call(shlex.split(cmdstr), stdout=f)
  
      def getColumnList(self):
          """Get a list of columns available in sdds file.
  
          Returns
          -------
          list
          """
          # command string
          cmdstr = "{} sddsquery -columnList {}".format(self.sif, self.filename)
  
          # add to command history
          self._addHistory(cmdstr)
  
          # run the command
          p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
          (output, err) = p.communicate()
          p_status = p.wait()
  
          # set the columnlist
          self.columnlist = [l.decode("utf-8") for l in output.splitlines()]
  
          return self.columnlist
  
      def getColumnValues(self, memory_threshold=100e6):
          """Get column values from file and return either
          as dataframe or return filename of converted data.
  
          Parameters
          ----------
          memory_threshold : float, optional
  
          Returns
          -------
          str | pandas.DataFrame
          """
          # make sure we have the column names
          if self.columnlist is None:
              self.getColumnList()
  
          # data can be large (bigger than memory)
          # use threshold to preve
          if os.path.getsize(self.filename) > memory_threshold:
              print(
                  "File is large, output redirected to file {}".format(
                      self.filename + "_columnvalues.dat"
                  )
              )
  
              # build command string
              cmdstr = "{} sdds2stream -col={} {} > {}".format(
                  self.sif,
                  ",".join(self.columnlist),
                  self.filename,
                  self.filename + "_columnvalues.dat",
              )
  
              # run command
              p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
              (output, err) = p.communicate()
              p_status = p.wait()
  
              return self.filename + "_columnvalues.dat"
  
          else:
              # build command string
              cmdstr = "{} sdds2stream -col={} {}".format(
                  self.sif, ",".join(self.columnlist), self.filename
              )
  
              # run command
              p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
              (output, err) = p.communicate()
              p_status = p.wait()
  
              # load file output into dataframe
              output = output.decode("utf-8")
              output = pd.read_csv(StringIO(output), names=self.columnlist, delim_whitespace=True)
  
              return output
  
      def getParameterList(self):
          """Get list of parameter names available in sdds file.
  
          Returns
          -------
          list
          """
          # command string
          cmdstr = "{} sddsquery -parameterList {}".format(self.sif, self.filename)
  
          # add to command history
          self._addHistory(cmdstr)
  
          # run command
          p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
          (output, err) = p.communicate()
          p_status = p.wait()
  
          # set the parameter list
          self.parameterlist = [l.decode("utf-8") for l in output.splitlines()]
  
          return self.parameterlist
  
      def getParameterValues(self):
          """Get parameter values from sdds file.
  
          Returns
          -------
          dataframe
          """
          cmdstr = "{} sddsprintout -parameters=* -spreadsheet {}".format(self.sif, self.filename)
          p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
          (output, err) = p.communicate()
          p_status = p.wait()
          df = pd.read_csv(
              StringIO(output.decode("utf-8")),
              error_bad_lines=False,
              delim_whitespace=True,
              skip_blank_lines=True,
              skiprows=1,
              names=["ParameterName", "ParameterValue"],
              index_col=False,
          )
  
          df = df.set_index("ParameterName", drop=True)
          df = pd.to_numeric(
              df.drop(["SVNVersion", "Stage", "PreviousElementName"], errors="ignore")[
                  "ParameterValue"
              ]
          )
          self.ParameterName = df
          return df
  
      def readParticleData(self, vary=False, long=False):
          """Read the particle tracking data.
  
          Parameters
          ----------
          vary : bool, optional
                                                                          adds a step column for distinguising between the
                                                                          different vary_element steps.
          Returns
          -------
          DataFrame (pandas or dask)
          """
          # if vary command was used in ele command file
          # convert data and add step column
          if vary:
              self.process_scan()
  
          if long:
              self.addCommand(
                  "sddsprocess",
                  define='column,delta,"p pCentral / 1 - " {}'.format(self.filename),
                  outfile="{}_longitudinal.{}".format(
                      self.filename.split(".")[0], self.filename.split(".")[1]
                  ),
              )
  
              # run command
              self.runCommand()
  
              # cmdstr = "{} sddsprocess -define=column,step,Step {} {}_processed.{}".format(
              #    self.sif, self.filename, *self.filename.split(".")
              # )
              # print(cmdstr)
              # p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
              newfilename = "{}_longitudinal.{}".format(*self.filename.split("."))
              print("Warning - auto filename set")
              print("Changed from {} to {}".format(self.filename, newfilename))
  
              # update the filename
              self.filename = newfilename
  
          self.getColumnList()
  
          # if the file type is binary convert it first
          if self.filetype == 0:
              self.convert(outfile=None)
  
          # try to load the column data
          df = self.getColumnValues()
  
          # check if data or filename is returned
          if isinstance(df, str):
              # if filename was returned lazy load the data using dask dataframe
              data = dd.read_csv(df, delimiter=" ", names=self.columnlist, header=None)
          else:
              data = df
  
          # if vary commmand was used group the data and add Turn column
          if vary:
              grouped = data.groupby(by="step")
  
              def f(group):
                  return group.join(
                      pd.DataFrame({"Turn": group.groupby("particleID").cumcount() + 1})
                  )
  
              data = grouped.apply(f)
  
          else:
              data["Turn"] = data.groupby("particleID").cumcount() + 1
  
          # return data eager or lazy depending on the datasize
          return data
  
      def process_scan(self):
          """Process sdds file containig particle tracking data where
          *vary_element* was used in the command file/
          """
          # add command
          self.addCommand(
              "sddsprocess",
              define="column,step,Step {}".format(self.filename),
              outfile="{}_processed.{}".format(
                  self.filename.split(".")[0], self.filename.split(".")[1]
              ),
          )
  
          # run command
          self.runCommand()
  
          # cmdstr = "{} sddsprocess -define=column,step,Step {} {}_processed.{}".format(
          #    self.sif, self.filename, *self.filename.split(".")
          # )
          # print(cmdstr)
          # p = subp.Popen(cmdstr, stdout=subp.PIPE, shell=True)
          newfilename = "{}_processed.{}".format(*self.filename.split("."))
          print("Warning - auto filename set")
          print("Changed from {} to {}".format(self.filename, newfilename))
  
          # update the filename
          self.filename = newfilename
  
      def sddsplot_tunediagram(self, **kwargs):
  
          self.addCommand("sddsresdiag", file1="resdiag.sdds")
          self.runCommand()
  
          newkwargs = {
              "columnNames": "nux,nuy",
              "file": self.filename,
              "scale": "0,1,0,1",
              "graph": "sym,fill,vary=subtype",
              "order": "spect",
              "split": "col=x",
              "col": "nux,nuy",
              "file2": "resdiag.sdds",
              "sever": None,
          }
          newkwargs = {**newkwargs, **kwargs}
          self.addCommand("sddsplot", **newkwargs)
  
          self.runCommand()
  
      def sddsplot_base(self, **kwargs):
          """
          Basic sddsplot command.
          """
          sddscommand = SDDSCommand(self.sif)
          cmd = sddscommand.getCommand("sddsplot", file=self.filename, **kwargs)
  
          # add to command history
          self._addHistory(cmd)
  
          sddscommand.runCommand(cmd)
  
      def sddsplot(
          self,
          file=None,
          columnNames="x,xp",
          scale="0,0,0,0",
          graph="dot,vary",
          order="spectral",
          split="columnBin=particleID",
          **kwargs,
      ):
          """
          Quick standard multi particle tracking plotting
          of phase space - add columnNames to plot other
          variables.
  
          Colors are assigned per particle ID.
  
          Parameters:
          -----------
          """
          if file is not None:
              newkwargs = {
                  "file": file,
                  "columnNames": columnNames,
                  "scale": scale,
                  "graph": graph,
                  "order": order,
                  "split": split,
              }
  
              newkwargs = {**newkwargs, **kwargs}
              self.sddsplot_base(**newkwargs)
          else:
              print("File missing.")
  
      def sddsplot_fma(
          self,
          file=None,
          col="x,y",
          split="column=diffusionRate",
          graph="sym,vary=subtype,fill,scale=2,fill",
          order="spectral",
          **kwargs,
      ):
          newkwargs = {
              "file": file,
              "col": col,
              "graph": graph,
              "order": order,
              "split": split,
          }
          if file is None:
              newkwargs.pop("file")
          newkwargs = {**newkwargs, **kwargs}
          self.sddsplot_base(**newkwargs)
  
      def generate_scan_dataset(self, datasetdict):
          """
  		Generates a file called "temp_scan.sdds" containing columns of values
  		to be used by elegant to scan over using vary_element method.
  
  		Arguments:
  		----------
  		datadict: dict
  			dictionary where the keys are the column headers and values are list of values to scan over
  			Note: all dict values need to have the same length
  		
  		Example:
  		--------
  		>>> datasetdc = {
  			"Q1" : [1.205055,1.555550505],
  			"Q2" : [-1.45000,-1.800000000],\
  			"Q3D": [-2.02000,-2.020000000],\
  			"Q4D": [1.408000,1.4080000000],\
  			"Q3T": [0.000000,0.0000000000],\
  			"Q4T": [0.000000,0.0000000000],\
  			"Q5T": [0.000000,0.0000000000],\
  			"S1" : [0.000000,0.0000000000],\
  			"S2" : [0.000000,0.0000000000],\
  			"S3D": [0.000000,0.0000000000],\
  			"S4D": [0.000000,0.0000000000],\
  			"S3T": [0.000000,0.0000000000],\
  			"S4T": [0.000000,0.0000000000]
  		}
  		>>> sdds = SDDS(sif, "temp.sdds",0)
  		>>> sdds.generate_scan_dataset(datasetdc)
  
  		"""
          cmd = f"{self.sif}  sddsmakedataset {self.filename} "
  
          for k, v in datasetdict.items():
              cmd += f"-column={k},type=double -data=" + ",".join([str(vv) for vv in v]) + " "
  
          subp.run(cmd, check=True, shell=True)
          self.sdds_scan = "temp.sdds"
  