#检查一些ckpt内部情况
import torch

def get_param(checkpoint):

    if 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
    elif 'model' in checkpoint:
        state_dict = checkpoint['model']
    else:
        state_dict = checkpoint
    return state_dict    

def set_param(checkpoint, param):

    if 'state_dict' in checkpoint:
        checkpoint['state_dict'] = param
    elif 'model' in checkpoint:
        checkpoint['model'] = param
    else:
        checkpoint = param
    return checkpoint 


def load_checkpoint_with_missing_or_exsessive_keys_strict(new_checkpoint,target_checkpoint):
    state_dict = get_param(new_checkpoint)
    correct_dict = get_param(target_checkpoint).copy()


    for key in state_dict.keys():
        if correct_dict.pop(key, None) is None:
            print(f"Keys not found in model, which are excessive keys in checkpoint : {key}")
    # it will be initialized randomly

    state_dict = get_param(new_checkpoint)
    correct_dict = get_param(target_checkpoint).copy()
    for key in correct_dict.keys():
        if key not in state_dict:
            print(
                f"{key} not in loaded checkpoint"
            )
            state_dict.update({key: correct_dict[key]})
        elif state_dict[key].shape != correct_dict[key].shape:
            print(
                f"incorrect shape {key}:{state_dict[key].shape} vs {correct_dict[key].shape}"
            )
            state_dict.update({key: correct_dict[key]})

    # if we have more keys just discard them
    correct_dict = get_param(target_checkpoint)
    new_state_dict = dict()
    for key in state_dict.keys():
        if key in correct_dict.keys():
            new_state_dict.update({key: state_dict[key]})

    # target_checkpoint["state_dict"] = new_state_dict
    set_param(target_checkpoint, new_state_dict)
    return target_checkpoint


def load_checkpoint_with_missing_or_exsessive_keys(new_checkpoint,target_checkpoint):
    state_dict = get_param(new_checkpoint)
    correct_dict = get_param(target_checkpoint).copy()


    for key in state_dict.keys():
        if correct_dict.pop(key, None) is None:
            print(f"Keys not found in model, which are excessive keys in checkpoint : {key}")
    # it will be initialized randomly

    state_dict = get_param(new_checkpoint)
    correct_dict = get_param(target_checkpoint).copy()
    for key in correct_dict.keys():
        if key not in state_dict:
            print(
                f"{key} not in loaded checkpoint"
            )
            state_dict.update({key: correct_dict[key]})
        elif state_dict[key].shape != correct_dict[key].shape:
            a = state_dict[key]
            b = correct_dict[key]
            if a.ndim == b.ndim and a.shape[:-1] == b.shape[:-1]:
                print(
                f"incorrect shape {key}:{state_dict[key].shape} vs {correct_dict[key].shape}  but copy anyway"
                )
                min_last_dim = min(a.shape[-1], b.shape[-1])
                with torch.no_grad():
                    b[..., :min_last_dim].copy_(a[..., :min_last_dim])
            else:
                print(
                f"incorrect shape {key}:{state_dict[key].shape} vs {correct_dict[key].shape}"
                )        
            state_dict.update({key: correct_dict[key]})

    # if we have more keys just discard them
    correct_dict = get_param(target_checkpoint)
    new_state_dict = dict()
    for key in state_dict.keys():
        if key in correct_dict.keys():
            new_state_dict.update({key: state_dict[key]})

    # target_checkpoint["state_dict"] = new_state_dict
    set_param(target_checkpoint, new_state_dict)
    return target_checkpoint

root_path = '/home/data1/smn/SPFormer/'
if_strict = False
zero_path = root_path + 'run/epoch_0001.pth'
cfg_checkpoint = root_path + 'run/Teacher_SC1024_scan.pth'
cfg_checkpoint_ = root_path + 'run/SC512_pretrain.pth'


temp1 = torch.load(cfg_checkpoint)
temp2 = torch.load(zero_path)
if if_strict:
    target_checkpoint = load_checkpoint_with_missing_or_exsessive_keys_strict(temp1, temp2)
else:
    target_checkpoint = load_checkpoint_with_missing_or_exsessive_keys(temp1, temp2)


torch.save(target_checkpoint,cfg_checkpoint_)
print("done")


