import safetensors
import os
import torch
from safetensors import safe_open


path = '/home/rl/Downloads/output/checkpoint-4'
path = '/media/rl/HDD/data/multi_head_train_results/aloha_qwen2_vla/qwen2_vl_2B/qwen2_vl_only_folding_shirt_lora_ema_finetune_dit_h_4w_steps/checkpoint-30000'
def compare_lora_weights():
    ckpt = safe_open(os.path.join(path, 'adapter_model.safetensors'), framework='pt')
    ema_ckpt = safe_open(os.path.join(path, 'ema', 'adapter_model.safetensors'), framework='pt')

    for k in ckpt.keys():
        # print(f">>>>>>>>>>>>>>>>>>>>>>{k}<<<<<<<<<<<<<<<<<<<<<<<")
        print(k, torch.equal(ckpt.get_tensor(k),ema_ckpt.get_tensor(k)))

    pass

def compare_non_lora_weights():
    ckpt = torch.load(os.path.join(path, 'non_lora_trainables.bin'))
    try:
        ema_ckpt = torch.load(os.path.join(path, 'ema_non_lora_trainables.bin'))
    except Exception as e:
        print(e)
        ema_ckpt = torch.load(os.path.join(path, 'ema', 'non_lora_trainables.bin'))

    for k in ckpt.keys():
        # print(f">>>>>>>>>>>>>>>>>>>>>>{k}<<<<<<<<<<<<<<<<<<<<<<<")
        print(k, torch.equal(ckpt[k], ema_ckpt[k]))

    pass

def compare_zero_weights(tag='global_step30000'):
    ckpt = torch.load(os.path.join(path, tag, 'bf16_zero_pp_rank_6_mp_rank_00_optim_states.pt'), map_location=torch.device('cpu'))['optimizer_state_dict']
    ema_ckpt = torch.load(os.path.join(path, 'ema', tag, 'bf16_zero_pp_rank_6_mp_rank_00_optim_states.pt'), map_location=torch.device('cpu'))['optimizer_state_dict']
    print(ckpt.keys())
    for k in ckpt.keys():
        # print(f">>>>>>>>>>>>>>>>>>>>>>{k}<<<<<<<<<<<<<<<<<<<<<<<")
        print(k, torch.equal(ckpt[k], ema_ckpt[k]))

    pass

def compare_ema_weights():
    ckpt = torch.load(os.path.join(path, 'non_lora_trainables.bin'), map_location=torch.device('cpu'))
    ema_ckpt = torch.load(os.path.join(path, 'ema_weights_trainable.pth'), map_location=torch.device('cpu'))
    # print(len(ema_ckpt.keys()), len(ckpt.keys()))
    for k in ema_ckpt.keys():
        # print(f">>>>>>>>>>>>>>>>>>>>>>{k}<<<<<<<<<<<<<<<<<<<<<<<")
        if 'policy_head' in k:
            bool_matrix = ckpt[k] == ema_ckpt[k]
            false_indices = torch.where(bool_matrix == False)
            print(k, bool_matrix, false_indices)
            for i,j in zip(false_indices[0], false_indices[1]):
                print(ckpt[k].shape, ckpt[k][i][j].to(ema_ckpt[k].dtype).item(), ema_ckpt[k][i][j].item())
            break
        if k in ckpt.keys():
            print(k, ckpt[k].dtype, ema_ckpt[k].dtype,  torch.equal(ckpt[k].to(ema_ckpt[k].dtype), ema_ckpt[k]))
        else:
            print(f'no weights for {k} in ckpt')

    pass
def debug():
    state_dict = model.state_dict()
    ema_state_dict = self.ema.averaged_model.state_dict()
    for k in ema_state_dict.keys():
        print(k, state_dict[k].requires_grad, torch.equal(state_dict[k], ema_state_dict[k]))



def check_norm_stats():
    path = '/media/rl/HDD/data/multi_head_train_results/aloha_qwen2_vla/qwen2_vl_2B/qwen2_vl_calculate_norm_stats/dataset_stats.pkl'
    import pickle

    with open(path, 'rb') as f:
        stats = pickle.load(f)
    gripper = {}
    for k, v in stats.items():
        gripper[k] = {}
        for kk, vv in v.items():
            gripper[k][kk] = [vv[6], vv[13]]
    pass

if __name__ == '__main__':
    # compare_non_lora_weights()
    # compare_zero_weights()
    # compare_ema_weights()
    # ema_ckpt = torch.load(os.path.join("/home/rl/Downloads/output/checkpoint-2", 'ema_weights.pth'), map_location=torch.device('cpu'))
    # for k,v in ema_ckpt.items():
    #     if
    check_norm_stats()
