import pandas as pd
import datetime
import h5py
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.colors import Normalize

from sys import exit
import argparse
import os

# coding=utf-8

#energy线图
k = np.loadtxt(open("k1.csv","rb"),delimiter=",",skiprows=0)
w = np.loadtxt(open("w1.csv","rb"),delimiter=",",skiprows=0)
wk = np.loadtxt(open("wk1.csv","rb"),delimiter=",",skiprows=0)
print("k="+str(k.shape))
print("w="+str(w.shape))
print("wk="+str(wk.shape))
X, Y = np.meshgrid(k, w)



k = np.loadtxt(open("k1.csv","rb"),delimiter=",",skiprows=0)
w = np.loadtxt(open("w1.csv","rb"),delimiter=",",skiprows=0)
wk = np.loadtxt(open("wk1.csv","rb"),delimiter=",",skiprows=0)
print("k="+str(k.shape))
print("w="+str(w.shape))
print("wk="+str(wk.shape))
X, Y = np.meshgrid(k, w)

# plot forward plopagating waves
fig = plt.figure()
ax1 = fig.add_subplot(111)
contourdata = ax1.pcolormesh(X, Y, wk, cmap='plasma',norm=Normalize(vmin=-10, vmax=-1))
pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
plt.xlabel('$k$')
plt.ylabel('$w$')
pp.set_label('$\log E_x$')
plt.savefig("Ex(k,w).png")
plt.close()



# plot forward plopagating waves
wk2 = np.loadtxt(open("wk2.csv","rb"),delimiter=",",skiprows=0)
fig = plt.figure()
ax1 = fig.add_subplot(111)
contourdata = ax1.pcolormesh(X, Y, wk2, cmap='plasma',norm=Normalize(vmin=-10, vmax=-1))
pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
plt.xlabel('$k$')
plt.ylabel('$w$')
pp.set_label('$\log E_y$')
plt.savefig("Ey(k,w).png")
plt.close()

# plot forward plopagating waves
wk3 = np.loadtxt(open("wk3.csv","rb"),delimiter=",",skiprows=0)
fig = plt.figure()
ax1 = fig.add_subplot(111)
contourdata = ax1.pcolormesh(X, Y, wk3, cmap='plasma',norm=Normalize(vmin=-10, vmax=-1))
pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
plt.xlabel('$k$')
plt.ylabel('$w$')
pp.set_label('$\log E_z$')
plt.savefig("Ez(k,w).png")
plt.close()

# plot forward plopagating waves
wk4 = np.loadtxt(open("wk4.csv","rb"),delimiter=",",skiprows=0)
fig = plt.figure()
ax1 = fig.add_subplot(111)
contourdata = ax1.pcolormesh(X, Y, wk4, cmap='plasma')
pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
plt.xlabel('$k$')
plt.ylabel('$w$')
pp.set_label('$\log B_y$')
plt.savefig("By(k,w).png")
plt.close()

# plot forward plopagating waves
wk5 = np.loadtxt(open("wk5.csv","rb"),delimiter=",",skiprows=0)
fig = plt.figure()
ax1 = fig.add_subplot(111)
contourdata = ax1.pcolormesh(X, Y, wk5, cmap='plasma',norm=Normalize(vmin=-10, vmax=-1))
pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
plt.xlabel('$k$')
plt.ylabel('$w$')
pp.set_label('$\log B_z$')
plt.savefig("Bz(k,w).png")
plt.close()

# plot forward plopagating waves
b0=0.1
by=(10**wk4)
bz=(10**wk5)
bw = (by ** 2 + bz ** 2) / b0 ** 2

fig = plt.figure()
ax1 = fig.add_subplot(111)
contourdata = ax1.pcolormesh(X, Y,np.log10(bw), cmap='plasma')
pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
plt.xlabel('$k$')
plt.ylabel('$w$')
pp.set_label('$\log B_w$')
plt.savefig("Bw(k,w).png")
plt.close()

# import scipy.io as sio
# import numpy as np
#
# # # 将单个变量保存为mat文件, 同目录下就会有一个x.mat文件, 可以在matlab中打开了
# # x = [[1, 2, 3, 4], [5, 6, 7, 8]]
# # sio.savemat('x.mat', {'x': x})
# #
# # # 将多个变量保存为mat文件
# # a, b, c, d = 1, 2, 3, 4
# # sio.savemat('abcd.mat', {'a': a, 'b': b, 'c': c, 'd': d})
#
# # 读取mat文件
# abcd = sio.loadmat('k.mat')
# print(abcd)