#!/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-052_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.15_Picard', \
             '20200716-01-051_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.01_Picard', \
             '20200716-01-049_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.001_Picard', \
             '20200716-01-056_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1Rand_Picard', \
             '20200716-01-060_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1Rand_Picard', \
             '20200716-01-061_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1Rand_Picard']
             

labels = ['Initial uniform $d=0.15$ m', 'Initial uniform $d=0.01$ m', 'Initial uniform $d=0.001$ m', 'Random nonuniform initial $d$', \
          'Random nonuniform initial $d$', 'Random nonuniform initial $d$']

min_Ms = []
max_Ms = []

multi_d = []
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'][()]
    multi_d.append(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 - multi_d[0]
    normdiff = np.max(np.abs(ddiff), axis=(1,2))
    print("normdiff:",normdiff)
    f.close()
    d_diffs.append(normdiff)

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

import matplotlib.pyplot as plt

fig = plt.figure()
fig.patch.set_facecolor('w')
numRows=1
numCols=1

plt.subplot(numRows,numCols,1)
for j in range(1,len(filenames)):
    plt.semilogy(d_diffs[j], '.-', color=colors[j-1], label=labels[j-1])

plt.legend(fontsize=10,loc=0)
plt.xlabel('Iteration')
plt.title('Difference in $d$ compared to $d_{initial} = 0.1$ m')
plt.ylabel(r'max$_{\theta,\zeta} |d_j - d_j^{0.1}|$  [meters]')

plt.tight_layout()

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