content: # -*- coding: utf-8 -*-
  """
      femagtools.forcedens
      ~~~~~~~~~~~~~~~~~~~~
  
      Read Force Density Plot Files
  
  
  """
  import os
  import re
  import glob
  import numpy as np
  import logging
  
  logger = logging.getLogger('femagtools.forcedens')
  
  filename_pat = re.compile(r'^(\w+)_(\d{3}).PLT(\d+)')
  num_pat = re.compile(r'([+-]?\d+(?:\.\d+)?(?:[eE][+-]\d+)?)\s*')
  pos_pat = re.compile(r'^\s*POSITION\s*\[(\w+)\]')
  unit_pat = re.compile(r'\[([^\]]+)')
  
  
  def _readSections(f):
      """return list of PLT sections
  
      sections are surrounded by lines starting with '[***'
      or 2d arrays with 7 columns
  
      Args:
        param f (file) PLT file to be read
  
      Returns:
        list of sections
      """
  
      section = []
      for line in f:
          if line.startswith('[****') or pos_pat.match(line):
              if section:
                  if len(section) > 2 and section[1].startswith('Date'):
                      yield section[1:]
                  else:
                      yield section
                  if line.startswith('[****'):
                      section = []
                  else:
                      section = [line.strip()]
          else:
              section.append(line.strip())
      yield section
  
  
  class ForceDensity(object):
  
      def __init__(self):
          self.type = ''
          self.positions = []
          pass
  
      def __read_version(self, content):
          rec = content[0].split(' ')
          if len(rec) > 3:
              self.version = rec[3]
          else:
              self.version = rec[-1]
              
      def __read_project_filename(self, content):
          self.project = content[1].strip()
          
      def __read_nodes_and_mesh(self, content):
          self.nodes, self.elements, self.quality = \
              [float(r[0]) for r in [num_pat.findall(l)
                                     for l in content[:3]]]
          for l in content[3:]:
              m = re.match(r'\*+([^\*]+)\*+', l)
              if m:
                  self.type = m.group(1).strip()
                  return
              
      def __read_date_and_title(self, content):
          d = content[0].split(':')[1].strip().split()
          dd, MM, yy = d[0].split('.')
          hh, mm = ''.join(d[1:-1]).split('.')
          self.date = '{}-{}-{}T{:02}:{:02}'.format(
              yy, MM, dd, int(hh), int(mm))
  
          if len(content) > 6:
              self.title = content[2].strip() + ', ' + content[6].strip()
          else:
              self.title = content[2].strip()
              
          self.current = float(num_pat.findall(content[4])[0])
          
      def __read_filename(self, content):
          self.filename = content[0].split(':')[1].strip()
          
      def __read_position(self, content):
          d = dict(position=float(num_pat.findall(content[0])[-1]),
                   unit=unit_pat.findall(content[0].split()[1])[0])
          cols = content[2].split()
          labels = cols[::2] # either X, FN, FT, B_N, B_T
                             # or X FX FY B_X B_Y
          d['column_units'] = {k: u for k, u in zip(labels,
                                                    [unit_pat.findall(u)[0]
                                                     for u in cols[1::2]])}
          m = []
          for l in content[4:]:
              rec = l.split()[1:]
              if len(rec) == len(labels):
                  m.append([float(x) for x in rec])
          d.update({k: v for k, v in zip(labels, list(zip(*m)))})
  
          self.positions.append(d)
  
      def read(self, filename):
          with open(filename) as f:
              for s in _readSections(f.readlines()):
                  logger.debug('Section %s' % s[0:2])
                  if s[0].startswith('FEMAG'):
                      self.__read_version(s)
                  elif s[0].startswith('Project'):
                      self.__read_project_filename(s)
                  elif s[0].startswith('Number'):
                      self.__read_nodes_and_mesh(s)
                  elif s[0].startswith('File'):
                      self.__read_filename(s)
                  elif s[0].startswith('Date'):
                      self.__read_date_and_title(s)
                  elif s[0].startswith('POSITION'):
                      self.__read_position(s)
  
      def fft(self):
          """return FFT of FN"""
          import scipy.fftpack
          try:
              ntiles = int(360/self.positions[0]['X'][-1])
              FN = np.tile(
                  np.array([p['FN'][:-1] for p in self.positions[:-1]]),
                  (ntiles, ntiles))
          except AttributeError:
              return []
  
          N = FN.shape[0]
          fdn = scipy.fftpack.fft2(FN)
          dim = N//ntiles//2
          return np.abs(fdn)[1:dim, 1:dim]/N
      
      def items(self):
          return [(k, getattr(self, k)) for k in ('version',
                                                  'type',
                                                  'title',
                                                  'current',
                                                  'filename',
                                                  'date',
                                                  'positions')]
  
      def __str__(self):
          "return string format of this object"
          if self.type:
              return "\n".join([
                  'FEMAG {}: {}'.format(self.version, self.type),
                  'File: {}  {}'.format(self.filename, self.date)] +
                               ['{}: {}'.format(k, v)
                                for k, v in self.items()])
          return "{}"
  
  
  def read(filename):
      f = ForceDensity()
      f.read(filename)
      return f
  
  
  def readall(workdir='.'):
      """collect all recent PLT files
      returns list of ForceDensity objects
      """
      plt = dict()
      pltfiles = sorted(glob.glob(os.path.join(workdir, '*_*.PLT*')))
      base = os.path.basename(pltfiles[-1])
      lastserie = filename_pat.match(base).groups()[1]
      for p in pltfiles:
          base = os.path.basename(p)
          m = filename_pat.match(base)
          if m and lastserie == m.groups()[1]:
              model, i, k = m.groups()
              fdens = ForceDensity()
              fdens.read(p)
              logging.info("%s: %s", p, fdens.title)
              if model in plt:
                  plt[model].append(fdens)
              else:
                  plt[model] = [fdens]
      return plt
  
  
  if __name__ == "__main__":
      import matplotlib.pyplot as pl
      import femagtools.plot
      import sys
      if len(sys.argv) == 2:
          filename = sys.argv[1]
      else:
          filename = sys.stdin.readline().strip()
  
      fdens = read(filename)
  
      # Show the results
  
      title = '{}, Rotor position {}'.format(
          fdens.title, fdens.positions[0]['position'])
      pos = fdens.positions[0]['X']
      FT_FN = (fdens.positions[0]['FT'],
               fdens.positions[0]['FN'])
      femagtools.plot.forcedens(title, pos, FT_FN)    
      pl.show()
  
      title = 'Force Density Harmonics'
      femagtools.plot.forcedens_fft(title, fdens)
      pl.show()
       
  