#!/usr/bin/env python

print "Include any arguments to save a PDF"

import numpy as np
from scipy.io import netcdf
import sys

save_pdf = len(sys.argv) > 1

filenames = ['20200716-01-045_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Picard', \
             '20200716-01-046_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Anderson2', \
             '20200716-01-047_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Anderson3']

labels = ['Picard','Anderson2','Anderson3']

min_Ms = []
max_Ms = []

d_diffs = []

for j in range(len(filenames)):
    filename = filenames[j] + '/regcoil_out.test.nc'
    f = netcdf.netcdf_file(filename, 'r', mmap=False)
    min_Ms.append(f.variables['min_M'][()])
    max_Ms.append(f.variables['max_M'][()])
    d = f.variables['d'][()]
    nfp = f.variables['nfp'][()]
    theta = f.variables['theta_coil'][()]
    zeta = f.variables['zeta_coil'][()]
    dtheta = theta[1] - theta[0]
    dzeta = zeta[1] - zeta[0]
    ddiff = d[1:,:,:] - d[:-1,:,:]
    #ddiff = np.full(ddiff.shape, 3.14) # For testing
    ddiffsq = ddiff * ddiff
    normdiff = np.sqrt(nfp * dtheta * dzeta * np.sum(ddiffsq,axis=(1,2)) / (4*np.pi*np.pi))
    normdiff = np.max(np.abs(ddiff), axis=(1,2))
    print("normdiff:",normdiff)
    f.close()
    d_diffs.append(normdiff)

#exit(0)
#colors = ['r','y','g','c','b','k']
colors = ['r', 'y', 'b']

import matplotlib.pyplot as plt

fig = plt.figure(figsize=(12,5))
fig.patch.set_facecolor('w')
numRows=1
numCols=2

mu0 = 4 * np.pi * (1.0e-7)
Mtarget = 1.4 / mu0

markerSize=5
plt.subplot(numRows,numCols,1)
for j in range(len(filenames)):
    plt.semilogy((max_Ms[j] - Mtarget)/Mtarget, '^-', color=colors[j], label=labels[j]+r', max$(M)/M_t - 1$',ms=markerSize)
    plt.semilogy((Mtarget - min_Ms[j])/Mtarget, 'v:', color=colors[j], label=labels[j]+r', $1 - $min$(M)/M_t$',ms=markerSize)

plt.legend(fontsize=10,loc=0)
plt.xlabel('Iteration')
#plt.ylabel('min(|M|) or max(|M|)')
plt.title('Relative departure from target magnetization magnitude $M_t$')

#########################################3

plt.subplot(numRows,numCols,2)
for j in range(len(filenames)):
    plt.semilogy(d_diffs[j], '.-', color=colors[j], label=labels[j])
plt.legend(fontsize=10,loc=0)
plt.xlabel('Iteration')
plt.title('Difference in magnet thickness between successive iterates')
#plt.ylabel(r'$||d_{j+1} - d_{j}||$ [m]') 
plt.ylabel(r'max$_{\theta,\zeta}|d_{j+1} - d_{j}|$ [meters]') 

#plt.tight_layout()
plt.subplots_adjust(top=0.95, bottom=0.09, left=0.05, right=0.99, hspace=0.2, wspace=0.26)

fontsize=14
plt.figtext(0.001, 0.999, '(a)', va='top', fontsize=fontsize)
plt.figtext(0.5, 0.999, '(b)', va='top', fontsize=fontsize)

if save_pdf:
    print "Saving PDF"
    plt.savefig(__file__ + ".pdf")
else:
    plt.show()
