import time
import numpy as np 
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec 
import torch 
import torch.nn as nn 
from utils.data import DataWithNoisyAndMaxDist
from models.BRNNDist import BRNN as Model, Loss, FocalLoss
plt.switch_backend('agg')
plt.rcParams['figure.figsize'] = (16, 12)
plt.rcParams['figure.dpi'] = 150

def main(args):
    model_name = f"ckpt/BRNNdist/BRNNDist.TW6144.v5.pt" #保存和加载神经网络模型权重的文件路径
    data_tool = DataWithNoisyAndMaxDist(stride=1, n_length=6144, padlen=256, noise_prob = 0.4, std = 0.1)
    #data_tool = DataPnSnWithPolarType(file_name="h5data", stride=1, n_length=20480, padlen=4096)
    device = torch.device("cuda:1")
    model = Model()
    try :
        model.load_state_dict(torch.load(model_name, map_location=device))
    except:
        pass
    model.to(device)
    model.train()
    if hasattr(model, "encoder"): model.encoder.eval()
    if hasattr(model, "rnns"):    model.rnns.eval()
    if hasattr(model, "dist_head"): model.dist_head.eval()    
    for p in model.parameters():
        p.requires_grad = False
    for p in model.decoder.parameters():
        p.requires_grad = True
    lossfn = Loss() 
    #lossDist = torch.nn.CrossEntropyLoss(label_smoothing=0.05)  # 轻度平滑，边界更稳
    #lossDist.to(device)
    class_w =  torch.tensor([0.520, 0.512, 1.968], device=device)
    focal = FocalLoss(alpha=class_w, gamma=1.5).to(device)
    acc_time = 0 #记录训练的累计时间
    outloss = open(f"logdir/loss/BRNN-dist/BRNNDist.TW6144.v5.txt", "a") #记录训练过程中的loss
    #optim = torch.optim.Adam(model.parameters(), 1e-4, weight_decay=0e-3)
    optim = torch.optim.AdamW(model.decoder.parameters(), lr=1e-4, weight_decay=1e-4)

    for step in range(100000):
        st = time.perf_counter() #记录当前时间，用于计算每个训练步骤的执行时间
        a1, a2, a3 = data_tool.batch_data(32) #获取一个批次的地震波数据 #128  32
        #print(a1.shape)
        #print(a3.shape)
        #print(b3.shape)
        #a1.shape (128, 20480, 3)
        #a3.shape (128, 20480, 5)
        #b3.shape (32,20480)
        wave = torch.tensor(a1, dtype=torch.float).to(device) 
        wave = wave.permute(0, 2, 1)
        d = torch.tensor(a2, dtype=torch.float32).to(device) 
        d = d.permute(0, 2, 1)
        labeldist = torch.tensor(a3, dtype=torch.long).to(device) 
        oc, od, odp = model(wave)
        #loss = focal(od, labeldist) 
        loss = lossfn(oc, d)
        loss.backward()
        if loss.isnan():
            print("NAN error")
            optim.zero_grad()
            continue 
        optim.step() 
        optim.zero_grad()
        ls = loss.detach().cpu().numpy()
        ed = time.perf_counter() #记录当前时间，用于计算每个训练步骤的执行时间
        outloss.write(f"{step},{ed - st},{ls}\n") #训练步骤、执行时间和损失值写入损失记录文件
        outloss.flush()
        acc_time += ed - st            
        if step % 100 == 0: #每训练100步执行
            d = a2[0]
            labeldist = a3[0]
            p = oc.detach().cpu().numpy()[0]
            out = odp.argmax(axis=1)
            pred = out[0]
            torch.save(model.state_dict(), model_name) #保存当前模型的权重
            gs = gridspec.GridSpec(3, 1,hspace=0.3) 
            fig = plt.figure(figsize=(16, 16), dpi=100) 
            for i in range(3):
                ax = fig.add_subplot(gs[i, 0])
                if i == 0:ax.set_title(f"Label:{labeldist},Logits:{pred}", ha="left", va="bottom", x=0.00, y=1.01)
                if i == 1:
                    ax.set_title("Pg", ha="left", va="bottom", x=0.00, y=1.01)       
                if i == 2:ax.set_title("Sg", ha="left", va="bottom", x=0.00, y=1.01)
                #max_index = np.argmax(d[:, 1])
                
                ax.plot(d[:, i], alpha=0.5, c="b") # 绘制震相标签，用蓝色表示
                ax.plot(p[i, :], alpha=0.9, c="r") #绘制预测，用红色表示
                w = a1[0, :, i%3]
                ax.plot(w, alpha=0.3, c="k")
                #ax.set_xlim(max_index - 1000, max_index + 1000)
            plt.savefig("logdir/trainfig/BRNN-dist/BRNNDist.TW6144.v5.png") 
            print(f"训练用时：{acc_time:6.1f}, 迭代次数：{step+1}, Loss:{ls:6.1f}")
    print("done!")

import argparse
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="拾取连续波形")          
    parser.add_argument('-d', '--dist', default=200, type=int, help="输入连续波形")       
    parser.add_argument('-o', '--output', default="result/t1", help="输出文件名")      
    parser.add_argument('-m', '--model', default="lppn.model", help="模型文件lppnmodel")                                                            
    args = parser.parse_args()      
    main(args)