## This is the user-script for plotting using gleaner tool.

# Path to gleaner (Better use environment variable PYTHONPATH!)
glrPath = '/home/gk779/apes/gleaner'
import sys
import os

# Do not use X-server to create and save plot
import matplotlib
#matplotlib.use('Agg')

## Import all required modules
import matplotlib.ticker as mtick
import matplotlib.pyplot as mplt
sys.path.append(glrPath)
import gleaner
import logging

# font setting
from matplotlib import rc
font_size = 12
font_family = 'serif'
font_type = 'Times New Roman'
rc('text',usetex=True)
font = {'family':font_family,'%s'%font_family:font_type,'size':font_size}
rc('font',**font)

#axis without scientific notation
y_formatter = mtick.ScalarFormatter(useOffset=False)

## -------------------------------------------------------------------------- ##
logging.basicConfig(level=logging.INFO)
## -------------------------------------------------------------------------- ##
logging.info('Started creating plots ...') 
 
import glob
import re
from operator import itemgetter

resfiles = glob.glob('./tracking_nocharge/*line_p00000_*.res')

import sqlite3
# data base filename
dbname = 'track_ani.db'
# load database if exist else load tracking files and add to database
if os.path.isfile(dbname):
  os.remove(dbname)

print ('Processing data from tracking files')

ftimes = []
for rf in resfiles:
  m = re.search('tracking_nocharge/.*_t(?P<time>.*).res', rf)
  timestr = rf.split('_')[-1]
  ftimes.append((float(m.group('time')), timestr, rf))

sortfiles = sorted(ftimes, key=itemgetter(0))

# Load text, dump into a database with specific tabname to get columns later
for iFile in range(len(sortfiles)):
  timestr = sortfiles[iFile][1] 
  print ('Processing time file: *'+timestr)
  # Load text, dump into a database with specific tabname to get columns later
  sqlcon = gleaner.tracking_to_db(fname = ['tracking_nocharge/*line_p*'+timestr], \
                                  dbname=dbname, tabname='line'+str(iFile))

## -------------------------------------------------------------------------- ##
# Animation for all time steps
logging.info('Started creating animation ...')

# File name of the animation
anim_file = 'potential.mp4'
from matplotlib import animation
# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']
writer = Writer(fps=10, metadata=dict(artist='Me'))#, bitrate=1800)

fig = mplt.figure()
ax = fig.add_subplot(111)
mplt.grid(True)
#mplt.legend(loc=8, ncol=1,borderaxespad=0, \
#            prop={'size':font_size}).get_frame().set_lw(0.0)
mplt.xlabel('x (m)')
mplt.ylabel('Potential ($kV$)')
mplt.grid(True,which="major",ls="-")
ax.yaxis.set_major_formatter(y_formatter)

get_data_for_cols = ['coordX','potential_phy','analy_pot']

hline = []
for iplt in range(len(sortfiles)):
  [x, pot, pot_a] = gleaner.get_columns(sqlcon, tabname='line'+str(iplt), \
                               columns=get_data_for_cols)
  x, pot, pot_a = zip(*sorted(zip(x,pot,pot_a))) # sort of needed
  pot_scaled = []
  for ipot in range(len(pot)):
    pot_scaled.append(pot[ipot][0]*1e-3)
  timestr = sortfiles[iplt][1].split('.res')[0]
  if iplt == 0:
    hline.append(mplt.plot(x, pot_scaled, '-', color = 'b', label = 't0'))
  else:  
    hline.append(mplt.plot(x, pot_scaled, '-', color = 'b', label = str(timestr)))

anim = animation.ArtistAnimation(fig, hline, interval=50, repeat_delay=3000, \
                                 blit=True)

anim.save(anim_file, writer=writer)
#mplt.show(True)
logging.info('Animation created')  

logging.info('Plots created')
