## 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


import sqlite3
# data base filename
dbname = 'track.db'
# load database if exist else load tracking files and add to database
if os.path.isfile(dbname):
  print ('Processing data from existing database')
#  os.remove(dbname)
  sqlcon = sqlite3.connect(dbname)
else:
  print ('Processing data from tracking files')

  ftimes = []
  resfiles = glob.glob('./tracking_nocharge/*line_p00000_*.res')
  for rf in resfiles:
    m = re.search('tracking_nocharge/.*_t(?P<time>.*).res', rf)
    ftimes.append((float(m.group('time')), rf))
  
  sortfiles = sorted(ftimes, key=itemgetter(0))
  last_ftime = sortfiles[-1][1].split('_')[-1]
  print ('Processing time file: *'+last_ftime)

  # Load text, dump into a database with specific tabname to get columns later
  sqlcon = gleaner.tracking_to_db(fname = ['tracking_nocharge/*line_p*'+last_ftime], \
                                  dbname=dbname, tabname='line_nocharge')

  ftimes = []
  resfiles = glob.glob('./tracking_charge/*line_p00000_*.res')
  for rf in resfiles:
    m = re.search('tracking_charge/.*_t(?P<time>.*).res', rf)
    ftimes.append((float(m.group('time')), rf))
  
  sortfiles = sorted(ftimes, key=itemgetter(0))
  last_ftime = sortfiles[-1][1].split('_')[-1]
  print ('Processing time file: *'+last_ftime)

  # Load text, dump into a database with specific tabname to get columns later
  sqlcon = gleaner.tracking_to_db(fname = ['tracking_charge/*line_p*'+last_ftime], \
                                  dbname=dbname, tabname='line_charge')
## -------------------------------------------------------------------------- ##
print ('Potential over x:')
fig = mplt.figure()
ax = fig.add_subplot(111)
# Plot x, y ... at certain time step

dia_inner = 2.77e-3 #m
dia_outer = 203.2e-3 #m
rad_inner = dia_inner/2.0
rad_outer = dia_outer/2.0
ref_pot = 0.0 
pot_inner = 50e3 #V
pot_outer = 0.0
permit = 8.854187e-12 
QC = 20e-6 # C/m^3
Q0 = 0
# Parameter for analytical solution
import math
def analy_pot(x, Q): 
  fac = Q/(4.0*permit)
  r = x
  a = rad_inner
  b = rad_outer
  log_fac = (np.log(r)-np.log(b))/(np.log(a)-np.log(b))
  return fac*(np.square(b)-np.square(r)) \
        + (pot_inner-fac*(np.square(b)-np.square(a)))*log_fac

import numpy as np
x1 = np.linspace(rad_inner, rad_outer,25)
pot_a = analy_pot(x1, Q0)
mplt.plot(x1*1e3, pot_a*1e-3, 'x', color = 'r', label = 'Analytic, $\\rho_e=0$')

# Simulation result
get_data_for_cols = ['coordX','potential_phy']
[x, y] = gleaner.get_columns(sqlcon, tabname='line_nocharge', \
                             columns=get_data_for_cols)
x, y = zip(*sorted(zip(x,y))) # sort of needed
ynew = []
xnew = []
for ii in range(len(y)):
  ynew.append(y[ii][0]*1e-3)
  xnew.append(x[ii][0]*1e-3)
mplt.plot(x, ynew, '-', color = 'k', label = 'Simulation, $\\rho_e=0$')


pot_a = analy_pot(x1, QC)
mplt.plot(x1*1e3, pot_a*1e-3, 'o', color = 'r', label = 'Analytic, $\\rho_e=20$ $\mu C/m^3$')

[x, y] = gleaner.get_columns(sqlcon, tabname='line_charge', \
                             columns=get_data_for_cols)
x, y = zip(*sorted(zip(x,y))) # sort of needed
ynew = []
xnew = []
for ii in range(len(y)):
  ynew.append(y[ii][0]*1e-3)
  xnew.append(x[ii][0]*1e3)
mplt.plot(x, ynew, '--', color = 'k', label = 'Simulation, $\\rho_e=20$ $\mu C/m^3$')

# plot setting
mplt.legend(loc=1, ncol=1,borderaxespad=0, \
            prop={'size':font_size}).get_frame().set_lw(0.0)
mplt.xlabel('Distance from inner electrode r ($mm$)')
mplt.ylabel('Potential ($kV$)')
mplt.grid(True,which="major",ls="-")
ax.yaxis.set_major_formatter(y_formatter)

# save fig
figsize = [8,6]
fig = mplt.gcf()
fig.set_size_inches(figsize[0],figsize[1])
mplt.savefig('PotentialOverX.pdf', dpi=100, format='pdf', \
             bbox_inches="tight",interpolation=None)

## -------------------------------------------------------------------------- ##


logging.info('Plots created')
