#!/usr/bin/env python3

print("Add any arguments to save a PDF")

import matplotlib as mpl
import matplotlib.pyplot as plt
#from matplotlib import cm
import numpy as np
from scipy.io import netcdf
from scipy.interpolate import interp1d
import sys
import os

makePDF = False
if len(sys.argv) > 1:
    makePDF = True

#filename = 'regcoil_out.20200706-01-001-c09r00_withPorts_lambda1e-15_Picard_thetaZeta128_mpolNtor32_ns2_mgrid.nc'
#filename = 'regcoil_out.20200716-01-053_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Picard_1T.nc'
filenames = ['wout_c09r00_fixedBoundary_0.5T_vacuum_ns201.nc', \
                 'wout_free_boundary_from_regcoil_pm_20200716-01-069_noPorts.nc', \
                 'wout_free_boundary_from_regcoil_pm_20200716-01-070_ports.nc']

rmncs = []
zmnss = []
xms = []
xns = []
iotas = []

for j in range(len(filenames)):
    filename = filenames[j]
    f = netcdf.netcdf_file(filename,'r',mmap=False)
    nfp = f.variables['nfp'][()]
    rmncs.append(f.variables['rmnc'][()])
    zmnss.append(f.variables['zmns'][()])
    xms.append(f.variables['xm'][()])
    xns.append(f.variables['xn'][()])
    iotas.append(f.variables['iotaf'][()])
    f.close()
    print("Read file ", filename)

########################################################
# Now make plot of surfaces at given toroidal angle
########################################################

figureNum = 1
fig = plt.figure(figureNum,figsize=(10,4))
fig.patch.set_facecolor('white')

numRows = 1
numCols = 3

#linespecs = ['-r', '--g', ':b']
#linewidths = [2, 1.5, 1]

linespecs = ['-r', '-y', '-b']
linewidths = [2, 1.75, 0.7]

phis = [0, np.pi/3]
titles = ['0', '\pi/3']

ntheta = 200
theta = np.linspace(0,2*np.pi,ntheta)
for j in range(2): # Loop over sub-plots
    plt.subplot(numRows,numCols,j+1)
    phi = phis[j]
    for jwout in range(3): # Loop over the 3 wout files
        for js in range(0, 201, 40):
            R = np.zeros(ntheta)
            Z = np.zeros(ntheta)
            for imn in range(len(xms[jwout])):
                m = xms[jwout][imn]
                n = xns[jwout][imn]
                angle = m * theta - n * phi
                R += rmncs[jwout][js, imn] * np.cos(angle)
                Z += zmnss[jwout][js, imn] * np.sin(angle)
            plt.plot(R, Z, linespecs[jwout], lw=linewidths[jwout])

    plt.gca().set_aspect('equal',adjustable='box')
    #if whichPlot==2:
    #    plt.legend(fontsize=7, framealpha=1.0)
    plt.title(r'$\phi = ' + titles[j] + '$')
    plt.xlabel('R [meters]')
    plt.ylabel('Z [meters]')
    #plt.xlim([Rmin - dR * margin, Rmax + dR * margin])
    #plt.ylim([Zmin - dZ * margin, Zmax + dZ * margin])

plt.subplot(numRows,numCols,3)
s = np.linspace(0,1,201,endpoint=True)
labels=['Fixed-boundary target','Free boundary, no ports', 'Free boundary, with ports']
for jwout in range(3):
    plt.plot(s, iotas[jwout], linespecs[jwout], lw=linewidths[jwout], label=labels[jwout])
plt.legend(loc=0, fontsize=9.5)
plt.xlabel('Normalized toroidal flux $s$')
plt.title('Rotational transform $\iota$')

plt.tight_layout()

abc=['a','b','c']
for j in range(3):
    plt.figtext(0.015 + 0.33*j, 0.925, '('+abc[j]+')', fontsize=15)

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

