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

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

try:
  freqscale
except:
  freqscale=0.0

try:
  bothm
except:
  bothm=True

try:
  dpi
except:
  dpi = 300

try:
  nframes
except:
  nframes = 64

try:
  colorshift
except:
  colorshift=1

try:
  animate
except:
  animate=1

try:
  show
except:
  show=0

try:
  caption
except:
  caption=0

try:
  pixels
except:
  pixels=1048

try:
  bangle
except:
  bangle=30.0

try:
  distobs
except:
  distobs=220.0

try:
  xoffset
except:
  xoffset=0.0

try:
  yoffset
except:
  yoffset=0.0

try:
  pangle
except:
  pangle=0.0

try:
  figsize
except:
  figsize=5

#try:
#  rsurf
#except:
#  rsurf=1.0

sosh.nframes=nframes
sosh.dpi=dpi
sosh.icolshift=colorshift
sosh.capflag=caption

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

#(phi,theta)=sosh.image2sphere(xpixels=pixels,ypixels=pixels,bangle=bangle,distobs=distobs,pangle=pangle,xoffset=xoffset,yoffset=yoffset)
xx = np.linspace(0, 2*pixels-1, 2*pixels)/pixels*np.pi-np.pi
yy = np.linspace(0, pixels-1, pixels)/pixels*np.pi
phi,theta = np.meshgrid(xx,yy)
x=np.cos(theta)
(nx,ny)=x.shape

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

arrlist=sosh.arrlist_surf
lmaxlist=sosh.lmaxlist_surf
#varlist={}

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

interval = 1.0/25
maxsave = 0.0

#Animation function to call
def calcsumt(i):

  sumr=0.0
  sumt=0.0
  sump=0.0
  plotvar=sosh.plotvar
  
  for k in range(nmodes):
    (l,m,n)=q[k,:]
    signedm=m
    m=np.abs(m)
    (plm,dplm)=arrlist[m]
    (y,dy)=(plm[...,l-m],dplm[...,l-m])
    
    if plotvar in ['Vr','Vmag','Vsq']:
      ylm = y*np.exp(1.0j*signedm*phi)
    if plotvar in ['Vt', 'Vh', 'Vmag','Vsq']:
      dylmt=-np.sin(theta)*dy*np.exp(1.0j*signedm*phi)
    if plotvar in ['Vp', 'Vh', 'Vmag','Vsq']:
      dylmp2=1.0j*signedm*y*np.exp(1.0j*signedm*phi)/np.sin(theta)
    if (bothm and m != 0):
      signedm *= -1
      if plotvar in ['Vr','Vmag','Vsq']:
         ylm += y*np.exp(1.0j*signedm*phi)
      if plotvar in ['Vt', 'Vh', 'Vmag','Vsq']:
        dylmt += -np.sin(theta)*dy*np.exp(1.0j*signedm*phi)
      if plotvar in ['Vp', 'Vh', 'Vmag','Vsq']:
        dylmp2 += 1.0j*signedm*y*np.exp(1.0j*signedm*phi)/np.sin(theta)

    if freqscale==0:
      if plotvar in ['Vr','Vmag','Vsq']:
        sumr += freqlist[k]*rclist[k]*ylm*np.exp(-1.0j*2*np.pi*float(i)/float(nframes))*np.exp(-1.0j*phaselist[k])
      if plotvar in ['Vt', 'Vh', 'Vmag','Vsq']:
        sumt += freqlist[k]*hclist[k]*dylmt*np.exp(-1.0j*2*np.pi*float(i)/float(nframes))*np.exp(-1.0j*phaselist[k])    
      if plotvar in ['Vp', 'Vh', 'Vmag','Vsq']:
        sump += freqlist[k]*hclist[k]*dylmp2*np.exp(-1.0j*2*np.pi*float(i)/float(nframes))*np.exp(-1.0j*phaselist[k])

    else:
      if plotvar in ['Vr','Vmag','Vsq']:
        sumr += freqlist[k]*rclist[k]*ylm*np.exp(-1.0j*2*np.pi*freqlist[k]*freqscale*float(i)/float(nframes))*np.exp(-1.0j*phaselist[k])
      if plotvar in ['Vt', 'Vh', 'Vmag','Vsq']:
        sumt += freqlist[k]*hclist[k]*dylmt*np.exp(-1.0j*2*np.pi*freqlist[k]*freqscale*float(i)/float(nframes))*np.exp(-1.0j*phaselist[k])
      if plotvar in ['Vp', 'Vh', 'Vmag','Vsq']:
        sump += freqlist[k]*hclist[k]*dylmp2*np.exp(-1.0j*2*np.pi*freqlist[k]*freqscale*float(i)/float(nframes))*np.exp(-1.0j*phaselist[k])

#  if plotvar in ['Vh', 'Vmag', 'Vsq']:
#    vh2=np.square(sumt.real) + np.square(sump.real)
#  if plotvar in ['Vmag', 'Vsq']:
#    v2=np.square(sumr.real) + vh2

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

  return d

def sumanimate(i):

  d=calcsumt(i)
  global maxsave
  mn,mx=d.min(),d.max()
  maxabs=np.abs([mn,mx]).max()
  if (maxabs > maxsave):
    maxsave=maxabs
#    print("new max=%f"%maxabs, end=" ", flush=True)

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

sosh.animate=sumanimate
sosh.calcimage=calcsumt
sosh.isave=0

print("You may enter 'q' to quit or 'c' to change parameters.")
#instr=sosh.catchc("Enter number of modes: ",2)
file=sosh.catchc("Enter name of mode list file: ", "modelist.txt")
if (file == 'q'):
  sys.exit()
q=np.loadtxt(file,dtype=int)
nmodes=q.shape[0]
freqlist=np.zeros(nmodes)
rclist=np.zeros(nmodes)
hclist=np.zeros(nmodes)
count=0
lstr=''
mstr=''
nstr=''
pm='\u00B1'
cap3=''
cap3b=r'$\ell$ = {0}, $m$ = {1}, $n$ = {2}' '\n'
for i in range(nmodes):
  print("MODE #%i" % (i+1))

#  (l,m,n)=sosh.querylmn(nmodes)
#  if l == -1:
#    sys.exit()
  (l,m,n)=q[i,:]
  print('l = ',l,'m = ',m,'n = ',n)
  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(1)
    ntest = nmod[l==lmod]

# not needed if reading surface values from table
#  xir, xih = sosh.getradial(l,n)
#  indsurf = np.abs(sosh.rmesh-rsurf).argmin()
#
#  rclist.append(xir[indsurf])
#  hclist.append(xih[indsurf])
#  rc=1.0
#  hc=1.0

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

  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.\n")
  else:
    (plm,dplm)=arrlist[m]
    (y,dy)=(plm[...,l-m],dplm[...,l-m])
    print("Spherical harmonic retrieved.\n")

#  ylm=y*np.exp(1.0j*signedm*phi)
#  dylmt=-np.sin(theta)*dy*np.exp(1.0j*signedm*phi)
#  dylmp=1.0j*signedm*y*np.exp(1.0j*signedm*phi)

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

  count += 1

  if (bothm and m != 0):
#    signedm *= -1
#    ylm=y*np.exp(1.0j*signedm*phi)
#    dylmt=-np.sin(theta)*dy*np.exp(1.0j*signedm*phi)
#    dylmp=1.0j*signedm*y*np.exp(1.0j*signedm*phi)
#    rlist.append(ylm)
#    tlist.append(dylmt)
#    plist.append(dylmp/np.sin(theta))
#    rclist.append(rc)
#    hclist.append(hc)
#    freqlist.append(freq)
    count += 1
    lstr = lstr + str(l) + ', '
    mstr = mstr + pm + str(m) + ', '
    nstr = nstr + str(n) + ', '
    cap3 = cap3 + cap3b.format(str(l),pm+str(m),str(n))
  else:
    lstr = lstr + str(l) + ', '
    mstr = mstr + str(signedm) + ', '
    nstr = nstr + str(n) + ', '
    cap3 = cap3 + cap3b.format(l,signedm,n)

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

lstr=lstr.strip(', ')
mstr=mstr.strip(', ')
nstr=nstr.strip(', ')
cap3=cap3.strip('\n')

sumr=0.0
sumt=0.0
sump=0.0
rng=np.random.default_rng()
phaselist=np.zeros(nmodes)
#freqlist=np.asarray(freqlist)
#for i in range(count):
##  sumr += rclist[i]*rlist[i]*freqlist[i]
##  sumt += hclist[i]*tlist[i]*freqlist[i]
##  sump += hclist[i]*plist[i]*freqlist[i]
#  sumr += rlist[i]
#  sumt += tlist[i]
#  sump += plist[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]
  var=calcsumt(0)
  maxsave=var.max()

  if (sosh.capflag == 1):
    capblank=r'$\ell$ = {0}' '\n' r'$m$ = {1}' '\n' r'$n$ = {2}'
  elif (sosh.capflag == 2):
    capblank='\n' r'$\ell$ = {0};  $m$ = {1};  $n$ = {2}' '\n'
  elif (sosh.capflag == 3):
    capblank=cap3
  else:
    capblank=''
  caption=capblank.format(lstr,mstr,nstr)

  fig,im = sosh.drawfigure(var, caption=caption, fsize=figsize)
  if (sosh.isave != 0):
#    sosh.savefigure(ianimate=animate, label='_surfacesum')
    if (animate == 0):
      plt.imsave("test.png", var, cmap=sosh.colormap)
    else:
      print("Writing files...")
      for i in range(nframes):
        print(i, end=" ", flush=True)
        plt.imsave("png_out/test{:03d}.png".format(i), calcsumt(i), cmap=sosh.colormap, vmin=-maxsave, vmax=maxsave)
      print("done.")
  if (show != 0):
    if (animate != 0):
      anim = animation.FuncAnimation(fig, sumanimate, frames=nframes, interval=interval)
    plt.show()

  print()
  c=sosh.queryplotparms()
  if (c == 'q'):
    break
  fs=sosh.catchc("Enter new freqscale: ", freqscale)
  if (fs == 'q'):
    break
  else:
    freqscale=float(fs)
  print("Average frequency =",freqscale*freqlist.mean())
  nf=sosh.catchc("Enter new nframes: ", nframes)
  if (nf == 'q'):
    break
  else:
    nframes=int(nf)
    sosh.nframes=nframes
  ry=sosh.catchc("Randomize phase? ", 'n')
  if (ry == 'q'):
    break
  elif (ry == 'y'):
    phaselist=rng.random(nmodes)*2*np.pi
  elif (ry == 'z'):
    phaselist=np.zeros(nmodes)