import os
import torch
import logging
import random
import torch.optim as optim
from tqdm import tqdm
# from torch.utils.tensorboard import SummaryWriter

from common.utils import *
from common.opt import opts
from common.h36m_dataset import Human36mDataset
from common.Mydataset import Fusion

from model.SGraFormer import sgraformer

import matplotlib.pyplot as plt  
from mpl_toolkits.mplot3d import Axes3D  
import numpy as np  

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
CUDA_ID = [0]
device = torch.device("cuda")



def visualize_skeletons(input_2D, output_3D, gt_3D, idx=5, output_dir='./output'):  
    # Ensure the tensors are on the CPU and convert them to numpy arrays  
    input_2D = input_2D.cpu().numpy()  
    output_3D = output_3D.cpu().numpy()  
    gt_3D = gt_3D.cpu().numpy()  

    # Get the first action and first sample from the batch  
    input_sample = input_2D[idx, 0]  
    output_sample = output_3D[idx, 0]  
    gt_3D_sample = gt_3D[idx, 0]  

    print(f'\ninput_sample shape: {input_sample.shape}')  
    print(f'output_sample shape: {output_sample.shape}')  

    fig = plt.figure(figsize=(25, 5))  

    # Define the connections (bones) between joints  
    bones = [  
        (0, 1), (1, 2), (2, 3),  # Left leg  
        (0, 4), (4, 5), (5, 6),  # Right leg  
        (0, 7), (7, 8), (8, 9), (9, 10),  # Spine  
        (7, 11), (11, 12), (12, 13),  # Right arm  
        (7, 14), (14, 15), (15, 16)   # Left arm  
    ]  

    # Colors for different parts  
    bone_colors = {  
        "leg": 'green',  
        "spine": 'blue',  
        "arm": 'red'  
    }  

    # Function to get bone color based on index  
    def get_bone_color(start, end):  
        if (start in [1, 2, 3] or end in [1, 2, 3] or   
            start in [4, 5, 6] or end in [4, 5, 6]):  
            return bone_colors["leg"]  
        elif start in [7, 8, 9, 10] or end in [7, 8, 9, 10]:  
            return bone_colors["spine"]  
        else:  
            return bone_colors["arm"]  

    # Plotting 2D skeletons from different angles  
    for i in range(4):  
        ax = fig.add_subplot(1, 7, i + 1)  
        ax.set_title(f'2D angle {i+1}')  
        ax.scatter(input_sample[i, :, 0], input_sample[i, :, 1], color='blue')  

        # Draw the bones  
        for start, end in bones:  
            bone_color = get_bone_color(start, end)  
            ax.plot([input_sample[i, start, 0], input_sample[i, end, 0]],  
                    [input_sample[i, start, 1], input_sample[i, end, 1]], color=bone_color)  

        ax.set_xlabel('X')  
        ax.set_ylabel('Y')  
        ax.set_xlim(np.min(input_sample[:, :, 0]) - 1, np.max(input_sample[:, :, 0]) + 1)  
        ax.set_ylim(np.min(input_sample[:, :, 1]) - 1, np.max(input_sample[:, :, 1]) + 1)  
        ax.grid()  

    # Plotting predicted 3D skeleton  
    ax = fig.add_subplot(1, 7, 5, projection='3d')  
    ax.set_title('3D Predicted Skeleton')  
    ax.scatter(output_sample[:, 0], output_sample[:, 1], output_sample[:, 2], color='red', label='Predicted')  

    # Draw the bones in 3D for output_sample  
    for start, end in bones:  
        bone_color = get_bone_color(start, end)  
        ax.plot([output_sample[start, 0], output_sample[end, 0]],  
                [output_sample[start, 1], output_sample[end, 1]],  
                [output_sample[start, 2], output_sample[end, 2]], color=bone_color)  

    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('Z')  
    ax.set_xlim(np.min(output_sample[:, 0]) - 1, np.max(output_sample[:, 0]) + 1)  
    ax.set_ylim(np.min(output_sample[:, 1]) - 1, np.max(output_sample[:, 1]) + 1)  
    ax.set_zlim(np.min(output_sample[:, 2]) - 1, np.max(output_sample[:, 2]) + 1)  
    ax.legend()  

    # Plotting ground truth 3D skeleton  
    ax = fig.add_subplot(1, 7, 6, projection='3d')  
    ax.set_title('3D Ground Truth Skeleton')  
    ax.scatter(gt_3D_sample[:, 0], gt_3D_sample[:, 1], gt_3D_sample[:, 2], color='blue', label='Ground Truth')  

    # Draw the bones in 3D for gt_3D_sample  
    for start, end in bones:  
        bone_color = get_bone_color(start, end)  
        ax.plot([gt_3D_sample[start, 0], gt_3D_sample[end, 0]],  
                [gt_3D_sample[start, 1], gt_3D_sample[end, 1]],  
                [gt_3D_sample[start, 2], gt_3D_sample[end, 2]], color=bone_color, linestyle='--')  

    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('Z')  
    ax.set_xlim(np.min(gt_3D_sample[:, 0]) - 1, np.max(gt_3D_sample[:, 0]) + 1)  
    ax.set_ylim(np.min(gt_3D_sample[:, 1]) - 1, np.max(gt_3D_sample[:, 1]) + 1)  
    ax.set_zlim(np.min(gt_3D_sample[:, 2]) - 1, np.max(gt_3D_sample[:, 2]) + 1)  
    ax.legend()  

    plt.grid()  

    # Save the figure  
    plt.tight_layout()  
    plt.savefig(f'{output_dir}/skeletons_visualization.png')  
    plt.show()  

def val(opt, actions, val_loader, model):
    with torch.no_grad():
        return step('test', opt, actions, val_loader, model)

def step(split, opt, actions, dataLoader, model, optimizer=None, epoch=None, writer=None, adaptive_weight=None):
    loss_all = {'loss': AccumLoss()}
    action_error_sum = define_error_list(actions)


    model.eval()

    TQDM = tqdm(enumerate(dataLoader), total=len(dataLoader), ncols=100)
    for i, data in TQDM:
        batch_cam, gt_3D, input_2D, action, subject, scale, bb_box, start, end, hops = data

        [input_2D, gt_3D, batch_cam, scale, bb_box, hops] = get_varialbe(split, [input_2D, gt_3D, batch_cam, scale, bb_box, hops])

        # print("\n======> input_2D: ", input_2D.shape)
        # print("======> gt_3D: ", gt_3D.shape)


        if split == 'train':
            output_3D = model(input_2D, hops)
        elif split == 'test':
            input_2D, output_3D = input_augmentation(input_2D, hops, model)

        out_target = gt_3D.clone()
        out_target[:, :, 0] = 0

        # print("======> output_3D: ", output_3D.shape)
        # visualize_skeletons(input_2D, output_3D, gt_3D)  

        if output_3D.shape[1] != 1:
            output_3D = output_3D[:, opt.pad].unsqueeze(1)
        output_3D[:, :, 1:, :] -= output_3D[:, :, :1, :]
        output_3D[:, :, 0, :] = 0
        action_error_sum = test_calculation(output_3D, out_target, action, action_error_sum, opt.dataset, subject)

        p1, p2 = print_error(opt.dataset, action_error_sum, opt.train)
        # print("======> p1, p2: ", p1, p2)

    if split == 'train':
        return loss_all['loss'].avg
    elif split == 'test':
        p1, p2 = print_error(opt.dataset, action_error_sum, opt.train)
        return p1, p2


def input_augmentation(input_2D, hops, model):
    input_2D_non_flip = input_2D[:, 0]
    output_3D_non_flip = model(input_2D_non_flip, hops)

    # print("======> input_2D_non_flip: ", input_2D_non_flip.shape)
    # print("======> output_3D_non_flip: ", output_3D_non_flip.shape)
    # visualize_skeletons(input_2D_non_flip, output_3D_non_flip)  
    
    return input_2D_non_flip, output_3D_non_flip




if __name__ == '__main__':
    opt = opts().parse()
    root_path = opt.root_path
    opt.manualSeed = 1
    random.seed(opt.manualSeed)
    torch.manual_seed(opt.manualSeed)


    root_path = opt.root_path
    dataset_path = root_path + 'data_3d_' + opt.dataset + '.npz'

    dataset = Human36mDataset(dataset_path, opt)
    actions = define_actions(opt.actions)

    train_data = Fusion(opt=opt, train=True, dataset=dataset, root_path=root_path)
    train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=opt.batch_size,
                                                    shuffle=True, num_workers=int(opt.workers), pin_memory=True)


    test_data = Fusion(opt=opt, train=False, dataset=dataset, root_path=root_path)
    test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=opt.batch_size,
                                                  shuffle=False, num_workers=int(opt.workers), pin_memory=True)

    model = sgraformer(num_frame=opt.frames, num_joints=17, in_chans=2, embed_dim_ratio=32, depth=4,
                      num_heads=8, mlp_ratio=2., qkv_bias=True, qk_scale=None, drop_path_rate=0.1)
    # model = FuseModel()

    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = torch.nn.DataParallel(model, device_ids=CUDA_ID).to(device)
    model = model.to(device)

    model_dict = model.state_dict()

    model_path = '/home/zlt/Documents/SGraFormer-master/checkpoint/epoch_50.pth'

    pre_dict = torch.load(model_path)

    model_dict = model.state_dict()
    state_dict = {k: v for k, v in pre_dict.items() if k in model_dict.keys()}
    model_dict.update(state_dict)
    model.load_state_dict(model_dict)

    all_param = []
    lr = opt.lr
    all_param += list(model.parameters())

    optimizer = optim.AdamW(all_param, lr=lr, weight_decay=0.1)

    ## tensorboard
    # writer = SummaryWriter("runs/nin")
    writer = None
    flag = 0



    p1, p2 = val(opt, actions, test_dataloader, model)


    print('p1: %.2f, p2: %.2f' % (p1, p2))
 
