import math
import glob
import numpy as np
from itertools import *

seeder_path = $!seeder_path!$
length = $!seeder_length!$
dx = $!seeder_dx!$
bc_type = $!seeder_bc_type!$
params = [bc_type]

plot  = []

show_plot = True
font_family= 'serif'
font_size = 12
set_majorgrid = True
figsize = [5,4]
markers = ['r-o','b-s','k-*','g-d','k-v','kx','k-p','k-^','k+']

def computeL2Norm(infiles, mus_file, incol, reference, refcol, fileprefix, computeSlope):
  error_file = fileprefix+'_l2Norm.txt'
  slope_file = fileprefix+'_slope.txt'
  #print(fileprefix)
  fileID = open(error_file,'w')
  for iFile in range(len(infiles)):
    #print(length/dx[iFile],round(length/dx[iFile]), int(length/dx[iFile]))
    dataFile = glob.glob(infiles[iFile]+mus_file)
    dataFile.sort()
    Ace_mus = np.array([])
    Met_mus = np.array([])
    Air_mus = np.array([])
    Ace_ode = np.array([])
    Met_ode = np.array([])
    Air_ode = np.array([])
    for iFile_2 in range(len(dataFile)):
      data = np.genfromtxt(dataFile[iFile_2], dtype = 'float')
      if data.size == data.shape[0]:
        data.resize(1,data.size)
      Ace_mus = np.append(Ace_mus,data[:, incol[0]])
      Met_mus = np.append(Met_mus,data[:, incol[1]])
      Air_mus = np.append(Air_mus,data[:, incol[2]])

    data = np.genfromtxt(infiles[iFile]+reference, dtype = 'float')
    Ace_ode = np.append(Ace_ode,data[:, refcol[0]])
    Met_ode = np.append(Met_ode,data[:, refcol[1]])
    Air_ode = np.append(Air_ode,data[:, refcol[2]])

    #print('orig size ', len(Ace_mus), len(Ace_ode))

    if 'moments' in fileprefix:
      slope = 2
      last = len(Ace_mus)-1
      Ace_mus = np.delete(Ace_mus,[0,last])
      Met_mus = np.delete(Met_mus,[0,last])
      Air_mus = np.delete(Air_mus,[0,last])
      
      Ace_ode = np.delete(Ace_ode,[0,last])
      Met_ode = np.delete(Met_ode,[0,last])
      Air_ode = np.delete(Air_ode,[0,last])

    if 'BounceBack' in fileprefix:
      slope = 1
      Ace_ode_tmp = np.array([])
      Met_ode_tmp = np.array([])
      Air_ode_tmp = np.array([])
      for i in range(len(Ace_ode)-1):
        Ace_ode_tmp = np.append(Ace_ode_tmp, (Ace_ode[i]+Ace_ode[i+1])/2.0)
        Met_ode_tmp = np.append(Met_ode_tmp, (Met_ode[i]+Met_ode[i+1])/2.0)
        Air_ode_tmp = np.append(Air_ode_tmp, (Air_ode[i]+Air_ode[i+1])/2.0)

      Ace_ode = Ace_ode_tmp
      Met_ode = Met_ode_tmp
      Air_ode = Air_ode_tmp

    nElems = len(Ace_mus)

    #Ace_error = np.linalg.norm(Ace_mus-Ace_ode)/nElems
    #Met_error = np.linalg.norm(Met_mus-Met_ode)/nElems
    #Air_error = np.linalg.norm(Air_mus-Air_ode)/nElems
    Ace_error = np.linalg.norm(Ace_mus-Ace_ode)/np.linalg.norm(Ace_ode)
    Met_error = np.linalg.norm(Met_mus-Met_ode)/np.linalg.norm(Met_ode)
    Air_error = np.linalg.norm(Air_mus-Air_ode)/np.linalg.norm(Air_ode)

    if computeSlope:
      if iFile == 0:
        # constant for linear slope
        c1_Ace = Ace_error/dx[iFile]
        c1_Met = Met_error/dx[iFile]
        c1_Air = Air_error/dx[iFile]
        # constant for quadratic slope
        c2_Ace = c1_Ace/dx[iFile]
        c2_Met = c1_Met/dx[iFile]
        c2_Air = c1_Air/dx[iFile]

        slopeID = open(slope_file, 'w')
        #print(dx[iFile], Air_error, c1, c2)
      slopeID.write( str(dx[iFile]/dx[0])+'\t'\
                    +str(c1_Ace*dx[iFile])+'\t'+str(c2_Ace*dx[iFile]*dx[iFile])+'\t'\
                    +str(c1_Met*dx[iFile])+'\t'+str(c2_Met*dx[iFile]*dx[iFile])+'\t'\
                    +str(c1_Air*dx[iFile])+'\t'+str(c2_Air*dx[iFile]*dx[iFile])+'\t'\
                    +'\n' )

    error = str(dx[iFile]/dx[0])+'\t'+str(round(length/dx[iFile]))+'\t'\
           +str(Ace_error)+'\t'+str(Met_error)+'\t'+str(Air_error)+'\n'
    fileID.write(error)
  fileID.close()
  if computeSlope:
    slopeID.close()
  return error_file


labels = []
filedict=dict()
filenames = []
for ipar in product(*params):
  pattern = 'bc_type'+str(ipar[0])
  filename = str(ipar[0]) 
  labels.append('bc\_type='+str(ipar[0]))
  filenames.append(filename)
  for iFile in seeder_path:
    if pattern in iFile:
      #print (filename, iFile)
      if filename not in filedict:
        filedict[filename] = [iFile]
      else:
        filedict[filename].append(iFile)


#print (filedict[filenames[0]])
#print (pattern)
#print (filenames)
#print (infiles[0])
#print ('moments/mus_omega0.8' in infiles[0])
error_files = []
slope_files = []
for iFile in range(len(filenames)):
  mus_file = '/tracking/stefantube_mix_line_p*_t*.res'
  ref_file = '/maxwell_stefan_reference.res'
  computeSlope = True
  if iFile == 0:
    computeSlope = True
  error_files.append(computeL2Norm(filedict[filenames[iFile]], mus_file, [3,4,5], \
                ref_file, [1,2,3], filenames[iFile], computeSlope))
  slope_files.append(filenames[iFile]+'_slope.txt')              


plot.append(dict(kind = 'xy', 
            data=[slope_files[0]],
            row=[1,-1], col=[1,7],dtype='float',ls='r', 
            #label = 'linear',
            xlogscale = True, 
            ylogscale = True,
            startplot = True))
plot.append(dict(kind = 'xy', 
            data=[slope_files[1]],
            row=[1,-1], col=[1,6],dtype='float',ls='b',
            xlogscale = True, 
            ylogscale = True,
            #label='quadratic'
            ))

for iFile in range(len(filenames)):
  sPlot = False
  ePlot = False
  if iFile == 0:
    sPlot = False
  if iFile == (len(filenames)-1):
    ePlot = True

  plot.append(dict(kind = 'xy', 
              data=[error_files[iFile]],
              row=[1,-1], col=[1,5],dtype='float', 
              startplot = sPlot, endplot = ePlot,
              #xlabel = "dx", 
              #ylabel = "error",
              label = labels[iFile], ls = markers[iFile], 
              #xmin=10, xmax=1e4,
              legend=dict(loc=6),
              xlogscale = True, 
              ylogscale = True,
              #label = '', ls = markers[ifile],
              #ls = '-o', 
              format = 'png', figname = 'molefraction_l2normError_differentBC', dpi=100 ))


