#!/usr/bin/env python

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import inspect, os
from scipy.io import netcdf
import sys

print("Add any argument to save a PDF")
makePDF = False
if len(sys.argv) > 1:
    makePDF = True

filenames = ['regcoil_out.20200716-01-045_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Picard.nc', \
                 'regcoil_out.20200716-01-059_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Picard_ports.nc']

ds = []
for j in range(2):
    filename = filenames[j]
    f = netcdf.netcdf_file(filename,'r',mmap=False)
# We use 'lambdas' instead of 'lambda' to avoid conflict with python's keyword lambda.
    lambdas = f.variables['lambda'][()]
    theta = f.variables['theta_coil'][()]
    zeta = f.variables['zeta_coil'][()]
    d = f.variables['d'][()]
    ds.append(d[-1,:,:])
    f.close()
    print("Read data from file "+filename)

print("d.shape: ", d.shape)

#fig = plt.figure(figsize=(11,7))
fig = plt.figure(figsize=(8,3.4))
fig.patch.set_facecolor('w')
numRows = 1
numCols = 2

dmaxval = np.max(d,axis=(0,1,2))
dminval = 0.0
dticks = np.arange(0, 0.14, 0.02)
#Mcontours = np.linspace(0,2500000., 25)
dcontours = np.arange(0, 0.14, 0.005)

from mpl_toolkits.axes_grid1.inset_locator import inset_axes
labels=['WITHOUT ports','WITH ports']

for jplot in range(2):
    plt.subplot(numRows, numCols, jplot + 1)
    data = ds[jplot].transpose()
    cnt=plt.contourf(zeta, theta, data, dcontours)
    # This is the fix for the white lines between contour levels
    for c in cnt.collections:
        c.set_edgecolor("face")
    print("max(d):", np.max(data,axis=(0,1)), "  min(d):", np.min(data,axis=(0,1)))
    plt.ylabel(r'$\theta$')
    plt.xlabel(r'$\zeta$')
    plt.title(r'  Magnet thickness $d$ [meters] ' + labels[jplot])
    plt.colorbar(ticks=dticks)


#plt.tight_layout()
plt.subplots_adjust(left=0.06, bottom=0.14, right=0.99, top=0.93)
if makePDF:
    print("Saving PDF")
    plt.savefig(__file__+'.pdf')
else:
    plt.show()

