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

#*********Turn off devide zero error in numpy**********#
np.seterr(divide = 'ignore')

#*********parser for get filename and select plot**********#
parser = argparse.ArgumentParser(description='Data plotter for KEMPO1')
#---关键语句---#
parser.add_argument('filename', nargs='*', help='file names of data <e.g> parabolic_051.h5')
#---关键语句---#
parser.add_argument('-p', '--power', action='store_true', help="plot the temporal and spatial profile of Bw")
parser.add_argument('-d', '--dynamic-spectrum', action='store_true', help='plot dynamic spectrum at x=20, 60')
parser.add_argument('-v', '--velocity-dist', action='store_true', help='plot velocity distribution at x=20')
parser.add_argument('-r', '--resonant-current', action='store_true', help='plot -je, jb and jb/Bw')
args = parser.parse_args()

if not args.filename:
    exit("Error: Filename is missing. Please specify data file.\n <e.g.> python plot_data.py parabolic_05.h5")

"""
Set plot 
"""
is_power = args.power
is_dynamic = args.dynamic_spectrum
is_vdist = args.velocity_dist
is_res = args.resonant_current

#--- 将is_power is_dynamic is_vdist is_res置为true ---#
if not (is_power | is_dynamic | is_vdist | is_res):
    is_power = True
    is_dynamic = True
    is_vdist = True
    is_res = True

# print(is_power)
# print(is_dynamic)
# print(is_vdist)
# print(is_res)
# print(datetime.datetime.now())

for file in args.filename:
    # create directory to store results
    dirname = os.path.dirname(file)
    basename_wo_ext = basename_without_ext = os.path.splitext(os.path.basename(file))[0]
    directory = os.path.join(dirname, basename_without_ext+'_result')
    os.makedirs(directory, exist_ok=True)
    # print(directory)
    # print(datetime.datetime.now())

    # ******** 读取h5文件 *********#
    datafile = h5py.File(file,'r')
    # print(datafile)
    # print(datetime.datetime.now())

    """
    读取输入参数
    """
    parameters = datafile["/parameters"]
    # print(parameters)

    dx = parameters["dx"].size
    dt = parameters["dt"][0]
    cv = parameters["cv"][0]
    # print("dx="+str(dx))
    # print("dt="+str(dt))
    # print("cv="+str(cv))

    omega_c = parameters["omega_c"][0]
    qm = parameters["qm"][0]
    b0 = omega_c / qm  #轴向静磁场
    # print("omega_c="+str(omega_c))
    # print("qm="+str(qm))
    # print("b0="+str(b0))

    nx = np.int(parameters["nx"][()])
    ntime = np.int(parameters["ntime"][()])
    damping_nd = np.int(parameters["damping_nd"][()]) #?
    # print("nx="+str(nx))
    # print("ntime="+str(ntime))
    # print("damping_nd="+str(damping_nd))
    # print(datetime.datetime.now())

    """
    Temporal and spatial profile of wave power
    """
    if is_power:
        # Magnitude of forward Bw
        field = datafile["/field"]
        x = field["x_axis"][damping_nd-1:nx-damping_nd] #不包含damping区域
        t = field["t_axis"][0:ntime-1:8] #8个数一组，取每组的第一个
        by_fwd = field["by_fwd"][damping_nd-1:nx-damping_nd, 0:ntime-1:8]
        bz_fwd = field["bz_fwd"][damping_nd-1:nx-damping_nd, 0:ntime-1:8]
        bw = (by_fwd ** 2 + bz_fwd ** 2)/b0**2
        X, Y = np.meshgrid(x,t)    #一维向量转化为二维网格矩阵
        # print("x_axis="+str(x))
        # print("t="+str(t))
        # print("by_fwd="+str(by_fwd))
        # print("bz_fwd="+str(bz_fwd))
        # print("bw="+str(bw))
        # print("X="+str(X))
        # print("Y="+str(Y))
        # print(datetime.datetime.now())
        # print("bw="+str(bw))
        # print("log10(bw)="+str(np.log10(bw)))
        # print("log10(bw).T="+str(np.log10(bw).T))

        # plot forward plopagating waves
        fig = plt.figure()
        ax1 = fig.add_subplot(111)
        contourdata = ax1.pcolormesh(X, Y, 10 * np.log10(bw).T, cmap='plasma', norm=Normalize(vmin=-80, vmax=-20))
        pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
        plt.xlabel('$x c^{-1}\Omega_e$')
        plt.ylabel('$t \Omega_e$')
        pp.set_label('$20\log B_w/B_0$')
        plt.savefig(os.path.join(directory, "bw_forward.png"))
        plt.close()

        # Magnitude of backward Bw
        by_bwd = field["by_bwd"][damping_nd - 1:nx - damping_nd, 0:ntime - 1:8]
        bz_bwd = field["bz_bwd"][damping_nd - 1:nx - damping_nd, 0:ntime - 1:8]
        bw = (by_bwd ** 2 + bz_bwd ** 2) / b0 ** 2

        # plot backward plopagating waves
        fig = plt.figure()
        ax1 = fig.add_subplot(111)
        contourdata = ax1.pcolormesh(X, Y, 10 * np.log10(bw).T, cmap='plasma', norm=Normalize(vmin=-80, vmax=-20))
        pp = fig.colorbar(contourdata, ax=ax1, orientation='vertical')
        plt.xlabel('$x c^{-1}\Omega_e$')
        plt.ylabel('$t \Omega_e$')
        pp.set_label('$20\log B_w/B_0$')
        plt.savefig(os.path.join(directory, "bw_backward.png"))
        plt.close()