#!/usr/bin/env python3

import numpy as np
from scipy.io import netcdf

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

import sys

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

filename='/Users/mattland/Box Sync/work19/regcoil_out.20190611-01-025_c09r00_withPorts_lambda1e-15_Picard_thetaZeta128_mpolNtor32_ns2_mgrid.nc'

f=netcdf.netcdf_file(filename,'r',mmap=False)
ports_weight = f.variables["ports_weight"][()]
theta = f.variables["theta_coil"][()]
zeta = f.variables["zeta_coil"][()]
f.close()

zeta = np.append(np.append(zeta,zeta+2*np.pi/3),zeta+4*np.pi/3)
ports_weight = np.kron([1,1,1],ports_weight)

import matplotlib.pyplot as plt
import matplotlib.colors as colors

fig = plt.figure(figsize=(6,3))
fig.patch.set_facecolor("w")

contours = 10 ** np.arange(0, 3.5, 0.25)
plt.contourf(zeta,theta,ports_weight,contours, norm=colors.LogNorm())
plt.xlabel(r"Toroidal angle $\zeta$")
plt.ylabel(r"Poloidal angle $\theta$")
plt.colorbar()
plt.title('Weight $w$ for excluding port regions')

plt.tight_layout()

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