import torch
from .ema import ExponentialMovingAverage

def load_model_weights(model, ckpt_path, use_ema=True, device='cuda:0'):
    """
    Load weights of a model from a checkpoint file.

    Args:
        model (torch.nn.Module): The model to load weights into.
        ckpt_path (str): Path to the checkpoint file.
        use_ema (bool): Whether to use Exponential Moving Average (EMA) weights if available.
    """
    checkpoint = torch.load(ckpt_path,map_location={'cuda:0': str(device)})
    total_iter = checkpoint.get('total_it', 0)
    
    if "model_ema" in checkpoint and use_ema:
        ema_key = next(iter(checkpoint["model_ema"]))
        if ('module' in ema_key) or ('n_averaged' in ema_key):
            model = ExponentialMovingAverage(model, decay=1.0)
            
        model.load_state_dict(checkpoint["model_ema"], strict=True)
        if ('module' in ema_key) or ('n_averaged' in ema_key):
            model = model.module
            print(f'\nLoading EMA module model from {ckpt_path} with {total_iter} iterations')
        else:
            print(f'\nLoading EMA model from {ckpt_path} with {total_iter} iterations')
    else:
        model.load_state_dict(checkpoint['encoder'], strict=True)
        print(f'\nLoading model from {ckpt_path} with {total_iter} iterations')

    return total_iter