import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
import sosh
import sys
import os

nargs=len(sys.argv)
#print(sys.argv)
if (nargs > 1):
  for i in range(1, nargs):
    exec(sys.argv[i])

try:
  freqscale
except:
  freqscale=1.0/3000

try:
  dpi
except:
  dpi = 300

try:
  nframes
except:
  nframes = 64

try:
  animateflag
except:
  animateflag=1

try:
  rsurf
except:
  rsurf=1.0

try:
  pixels
except:
  pixels=1000

try:
  bangle
except:
  bangle=-30

try:
  figsize
except:
  figsize=5

# not needed if using table of surface values
sosh.loadmodel()

#(phi,theta)=sosh.image2sphere(xpixels=pixels,ypixels=pixels,bangle=bangle)
(r, theta) = sosh.image2rtheta(xpixels=pixels,ypixels=pixels)
x=np.cos(theta)
(nx,ny)=x.shape
if (rsurf < 1.0):
  newind = (r <= rsurf)
  r.mask = np.logical_not(newindex)

#modeparms=np.loadtxt('modes.average')
# use model values instead of measured values
modeparms=np.loadtxt('modes.model.txt')
lmod=modeparms[:,0]
nmod=modeparms[:,1]

arrlist={}
lmaxlist={}
varlist={}

rlist=[]
tlist=[]
plist=[]
freqlist=[]
amplist=[]
rclist=[]
hclist=[]

interval = 1.0/25

#Animation function to call
def sumanimate(i):
  sumr=0.0
  sumt=0.0
  sump=0.0
  for k in range(nmodes):
    sumr += freqlist[k]*rclist[k]*rlist[k]*np.exp(-1.0j*(2*np.pi*freqlist[k]*float(i) / np.float(nframes)))
    sumt += freqlist[k]*hclist[k]*tlist[k]*np.exp(-1.0j*(2*np.pi*freqlist[k]*float(i) / np.float(nframes)))
    sump += freqlist[k]*hclist[k]*plist[k]*np.exp(-1.0j*(2*np.pi*freqlist[k]*float(i) / np.float(nframes)))
  vh2=(np.square(sumt.real) + np.square(sump.real))
  v2=(np.square(sumr.real) + vh2)

  plotvar=sosh.plotvar
  if (plotvar == 'Vr'):
    d=sumr.real
  if (plotvar == 'Vt'):
    d=sumt.real
  if (plotvar == 'Vp'):
    d=sump.real
  if (plotvar == 'Vh'):
    d=np.sqrt(vh2)
  if (plotvar == 'Vmag'):
    d=np.sqrt(v2)
  if (plotvar == 'Vsq'):
    d=v2

  im.set_data(d)
  fig.canvas.draw()
  return

sosh.animate=sumanimate
sosh.isave=0

print("You may enter 'q' to quit or 'c' to change parameters.")
instr=sosh.catchc("Enter number of modes: ",2)
if (instr == 'q'):
  sys.exit()
nmodes=int(instr)
for i in range(nmodes):
  print("MODE #%i" % (i+1))

  (l,m,n)=sosh.querylmn()
  if l == -1:
    sys.exit()
  ntest = nmod[l==lmod]
  while (n not in ntest):
    print("That n was not modelled for that l. Modelled values are in the range %i to %i. Try again." \
          % (ntest.min(),ntest.max()))
    (l,m,n)=sosh.querylmn()
    ntest = nmod[l==lmod]

  signedm=m
  m=np.abs(m)
  if m not in lmaxlist.keys() or l > lmaxlist[m]:
    lmaxlist[m]=l
    plm=np.ma.array(np.zeros((nx,ny,l-m+1)))
    dplm=np.ma.array(np.zeros((nx,ny,l-m+1)))
    (y,dy)=sosh.setplm(l,m,x,plm,dplm)
    arrlist[m]=(plm,dplm)
    print("Spherical harmonic calculated and saved.")
  else:
    (plm,dplm)=arrlist[m]
    (y,dy)=(plm[...,l-m],dplm[...,l-m])
    print("Spherical harmonic retrieved.")

  if (m > 0):
    phi=(np.pi/4)/m
  else:
    phi=0.0
  ylm=y*np.exp(1.0j*signedm*phi)
  dylmt=-np.sin(theta)*dy*np.exp(1.0j*signedm*phi)
  dylmp=1.0j*m*y*np.exp(1.0j*signedm*phi)

  rlist.append(ylm)
  tlist.append(dylmt)
  plist.append(dylmp/np.sin(theta))

  xir, xih = sosh.getradial(l,n)
  fr=np.interp(r,sosh.rmesh,xir*np.sqrt(sosh.rho))
  fh=np.interp(r,sosh.rmesh,xih*np.sqrt(sosh.rho))

  ind = ((l==lmod) & (n==nmod))
  freq = float(modeparms[ind,2])*freqscale
#  rc = float(modeparms[ind,5])/1e8
#  hc = float(modeparms[ind,6])/1e8
  rclist.append(fr)
  hclist.append(fh)
  freqlist.append(freq)
#  print("Using vertical to horizontal ratio of %f, theoretical is %f." % (rc/hc, float(modeparms[ind,4])))
  print("Scaled frequency = %f." % (freq))

# if using measured mode parameters, this would 
# allow you to compute frequency as a function of m
# from the fitted a-coefficients
#  ai=np.append([0.0],modeparms[ind,12:18]/1000)
#  pols=sosh.apols(l,6)
#  fx=np.matmul(pols,ai)
#  freq=modeparms[ind,2]+fx
#  freqlist.append(freq[l+signedm]*freqscale)
#  amplist.append(modeparms[ind,3])

sumr=0.0
sumt=0.0
sump=0.0
for i in range(nmodes):
  sumr += rclist[i]*rlist[i]*freqlist[i]
  sumt += hclist[i]*tlist[i]*freqlist[i]
  sump += hclist[i]*plist[i]*freqlist[i]

vh2=(np.square(sumt.real) + np.square(sump.real))
v2=(np.square(sumr.real) + vh2)

varlist['Vr']=sumr.real
varlist['Vt']=sumt.real
varlist['Vp']=sump.real
varlist['Vh']=np.sqrt(vh2)
varlist['Vmag']=np.sqrt(v2)
varlist['Vsq']=v2

while True:

  var=varlist[sosh.plotvar]

  fig,im = sosh.drawfigure(var, fsize=figsize)
  if (sosh.isave != 0):
    sosh.savefigure(animateflag=animateflag)
  if (animateflag != 0):
    anim = animation.FuncAnimation(fig, sumanimate, frames=nframes, interval=interval)
  plt.show()

  print()
  c=sosh.queryplotparms()
  if (c == 'q'):
    break