#!/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

filename = '20200716-01-045_nthetanzetaPlasma128_nthetanzetaCoil129_mpolntor64_ns3_dinit0.1_Picard/regcoil_out.test.nc'

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'][()]
abs_M = f.variables['abs_M'][()]
Bnormal_total = f.variables['Bnormal_total'][()]
f.close()

Nplots = 5

print("Read data from file "+filename)
print("abs_M.shape: ", abs_M.shape)
print("d.shape: ", d.shape)

#fig = plt.figure(figsize=(11,7))
fig = plt.figure(figsize=(13,7))
fig.patch.set_facecolor('w')
numRows = 2
numCols = Nplots

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)

Mmaxval = np.max(abs_M, axis=(0,1,2,3))
Mminval = np.min(abs_M, axis=(0,1,2,3))
#Mticks = [1e4, 1e5, 1e6]
#Mcontours = np.exp(np.linspace(np.log(Mminval),np.log(Mmaxval), 25))
Mticks = np.arange(0, 2.9, 0.2)
Mcontours = np.arange(0, 2.9, 0.05)

from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

for jplot in range(Nplots):
    plt.subplot(numRows, numCols, jplot + 1)
    data = d[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$', labelpad=-2)
    plt.xlabel(r'$\zeta$')
    plt.title(r'$d$ [m]')
    if jplot == Nplots - 1:
        ax2 = plt.gca()
        cax = inset_axes(ax2,
                           width="5%",  # width = 5% of parent_bbox width
                           height="100%",  # height : 50%
                           loc='lower left',
                           bbox_to_anchor=(1.05, 0., 1, 1),
                           bbox_transform=ax2.transAxes,
                           borderpad=0,
                       )
        plt.colorbar(ticks=dticks, cax=cax)

    plt.subplot(numRows, numCols, jplot + 1 + numCols)
    data = abs_M[jplot, 0, :, :].transpose() / (1.0e6)
    #cnt=plt.contourf(zeta, theta, data, Mcontours, norm=matplotlib.colors.LogNorm(vmax=Mmaxval,vmin=Mminval))
    cnt=plt.contourf(zeta, theta, data, Mcontours)
    # This is the fix for the white lines between contour levels
    for c in cnt.collections:
        c.set_edgecolor("face")
    print("max M:", np.max(data, axis=(0,1)), "  min M:", np.min(data, axis=(0,1)))
    plt.ylabel(r'$\theta$', labelpad=-2)
    plt.xlabel(r'$\zeta$')
    plt.title('$M$ [MA/m]')
    if jplot == Nplots - 1:
        #left, bottom, width, height = plt.gca().get_position().bounds
        #cax = fig.add_axes([left+width+0.01, bottom, 0.03, height])
        #plt.colorbar(cax=cax)
        ax2 = plt.gca()
        cax = inset_axes(ax2,
                           width="5%",  # width = 5% of parent_bbox width
                           height="100%",  # height : 50%
                           loc='lower left',
                           bbox_to_anchor=(1.05, 0., 1, 1),
                           bbox_transform=ax2.transAxes,
                           borderpad=0,
                       )
        plt.colorbar(ticks=Mticks, cax=cax)


    plt.figtext(0.115 + 0.19*jplot, 0.995, 'Iteration '+str(jplot), ha='center', va='top', fontsize=14)

#plt.tight_layout()
plt.subplots_adjust(top=0.895, bottom=0.074, left=0.034, right=0.949, hspace=0.318, wspace=0.239)

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

