#!/usr/bin/env python
import matplotlib as mpl 
import matplotlib.pyplot as plt
#from mpl_toolkits.mplot3d import Axes3D
import numpy as np
#from numpy import arange,meshgrid,sqrt,zeros
import sys
import argparse
import array as arr
import primefac
#import math
program_name = sys.argv[0]
arguments = sys.argv[1:]
count = len(arguments)
if (count!=3):
	print 'Aufruf: Datei-Prefix X-Dimension Y-Dimension'
	sys.exit(0)
print arguments
fileprefix=arguments[0]
plane_id=fileprefix[1] # X, Y oder Z-Plane wird im Dateinamen fest gelegt
#filestr='EZ1_043'
inlist_x=list()
inlist_y=list()
inlist_z=list()
#z_array=arr.array('f')
file_x = open(fileprefix+'X.out', 'r')
file_y = open(fileprefix+'Y.out', 'r')
file_z = open(fileprefix+'Z.out', 'r')
instr_x =  file_x.readline()
instr_y =  file_y.readline()
instr_z =  file_z.readline()
while len(instr_x):
	for i in range(0, len(instr_x)/12):
		inlist_x.append(instr_x[i*12:i*12+12])
		inlist_y.append(instr_y[i*12:i*12+12])
		inlist_z.append(instr_z[i*12:i*12+12])
	instr_x =  file_x.readline()
	instr_y =  file_y.readline()
	instr_z =  file_z.readline()
#print(inlist_x)	
#print len(inlist_x)
#adim=int(np.sqrt(len(inlist_x)))
xdim=int(arguments[1])
ydim=int(arguments[2])
inlist_x_len=len(inlist_x)
if (xdim*ydim!=inlist_x_len):
	factors = list( primefac.primefac(inlist_x_len) )
	print 'Dimensionsangabe falsch: Matrix Groesse: ', inlist_x_len, factors
	sys.exit(0)

x=0
#print("Format:", adim)
xv=np.zeros((xdim, ydim),dtype=float )
yv=np.zeros((xdim, ydim),dtype=float )
zv=np.zeros((xdim, ydim),dtype=float )
mv=np.zeros((xdim, ydim),dtype=float )
for i in range(0, xdim):
	for k in range(0, ydim):
		xv[i,k]=float(inlist_x[x])
		yv[i,k]=float(inlist_y[x])
		zv[i,k]=float(inlist_z[x])
		mv[i,k]=np.sqrt(xv[i,k]**2+yv[i,k]**2+zv[i,k]**2)
		x=x+1
		#print(i, k)
#print(xv)		
#print(z_array)
# http://folk.uio.no/oyvindry/mat1110/2016/3dplot_inf1100.pdf
# https://www.labri.fr/perso/nrougier/teaching/matplotlib/
# https://stackoverflow.com/questions/9170838/surface-plots-in-matplotlib
# https://predictablynoisy.com/matplotlib/api/_as_gen/matplotlib.pyplot.quiver.html
# https://predictablynoisy.com/matplotlib/gallery/color/colormap_reference.html
# https://pypi.org/project/primefac/
#u,v = np.arange(0,17,1),np.arange(0,17,1)
#u,v = np.meshgrid(u,v)
#M = np.sqrt(u*u+v*v) # magnitude
print 'Xmax: ', np.max(xv), '  Xmin: ', np.min(xv)
print 'Ymax: ', np.max(yv), '  Ymin: ', np.min(yv)
print 'Zmax: ', np.max(zv), '  Zmin: ', np.min(zv)
print fileprefix[1]
x_max=np.max(np.absolute(xv))
y_max=np.max(np.absolute(yv))
z_max=np.max(np.absolute(zv))

if plane_id=='Y':
	z_xy=y_max/(np.sqrt(x_max**2+z_max**2))
elif plane_id=='X':
	z_xy=x_max/(np.sqrt(z_max**2+y_max**2))
else:
	z_xy=z_max/(np.sqrt(x_max**2+y_max**2))
print 'Z:XY: ', z_xy
ar_wd=0.002
if (z_xy>0.5):
	ar_wd=ar_wd*z_xy*2
if (ar_wd>0.015): ar_wd=0.015

plt.figure(figsize=(8,6))
plt.title(arguments[0]+'   '+plane_id+':plane = '+str(round(z_xy,2)))
#qq=plt.quiver(x,y,u,v,M,cmap=plt.cm.jet)
if plane_id=='Y':
	qq=plt.quiver(xv,zv,mv,cmap=plt.cm.jet,minlength=4,width=ar_wd)
elif plane_id=='X':
	qq=plt.quiver(zv,yv,mv,cmap=plt.cm.jet,minlength=4,width=ar_wd)
else:
	qq=plt.quiver(xv,yv,mv,cmap=plt.cm.jet,minlength=4,width=ar_wd)
plt.colorbar(qq, cmap=plt.cm.jet)
plt.show()