#!/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-039_nthetanzetaPlasma128_nthetanzetaCoil128_mpolntor32_ns3_multipleOf10lambdas/regcoil_out.test.nc'

specialLambdas = [1e-9, 1e-15, 1e-18]

lambda_many = []
chi2_B_many = []
chi2_M_many = []
max_B_many = []
max_M_many = []
max_lambda = 0

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

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

fig = plt.figure(figsize=(11,7))
fig.patch.set_facecolor('w')
numRows = 2
numCols = 3
#f, axarr = plt.subplots(numRows, numCols)

Mmaxval = 2541888.
Mminval = 200.0
Mticks=[1e3, 1e4, 1e5, 1e6, 1e7]
#Mcontours = np.linspace(0,2500000., 25)
Mcontours = np.exp(np.linspace(np.log(Mminval),np.log(Mmaxval), 25))

Bmaxval = 0.3383
#Bminval = 0.00001
Bminval = 0.000001
Bticks = [1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1e+0]
Bcontours = np.exp(np.linspace(np.log(Bminval),np.log(Bmaxval), 25))


for jlambda in range(3):
    lambdaval = specialLambdas[jlambda]
    j = None
    for j in range(len(lambdas)):
        if np.abs(lambdaval - lambdas[j]) < 1.0e-20 and lambdas[j] > 1e-20:
            index = j

    plt.subplot(numRows, numCols, jlambda + 1)
    cnt=plt.contourf(zeta, theta, abs_M[index, 0, :, :].transpose(), Mcontours, norm=matplotlib.colors.LogNorm(vmax=Mmaxval,vmin=Mminval))
    # 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(abs_M[index, 0, :, :], axis=(0,1)))
    plt.ylabel(r'$\theta$')
    plt.xlabel(r'$\zeta$')
    plt.title('M [A/m]')
    if True:
        #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)
        plt.colorbar(ticks=Mticks)

    plt.subplot(numRows, numCols, jlambda + 1 + numCols)
    data = np.abs(Bnormal_total[index, :, :].transpose())
    data[data<Bminval] = Bminval*1.00001
    cnt=plt.contourf(zeta, theta, data, Bcontours, norm=matplotlib.colors.LogNorm(vmax=Bmaxval,vmin=Bminval))
    # This is the fix for the white lines between contour levels
    for c in cnt.collections:
        c.set_edgecolor("face")
    print("max(B_n):", np.max(np.abs(Bnormal_total[index, :, :].transpose()),axis=(0,1)))
    plt.ylabel(r'$\theta$')
    plt.xlabel(r'$\zeta$')
    plt.title(r'$|B_n|$ [T]')
    if True:
        plt.colorbar(ticks=Bticks)

    plt.figtext(0.16 + 0.33*jlambda, 0.995, r'$\lambda=10^{' + str(int(np.log10(lambdaval))) + '}$ T$^2/$A$^2$', ha='center', va='top', fontsize=14)

#plt.tight_layout()
plt.subplots_adjust(top=0.895, bottom=0.074, left=0.039, right=0.989, hspace=0.318, wspace=0.154)

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

