import glob
import hashlib

import torch
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner


PFX = "optimizer.state.fp32_param."


def load(name, prefix=PFX):
    path = glob.glob(f"ckpt/{name}/*/miles")[0]
    reader = dcp.FileSystemReader(path)
    metadata = reader.read_metadata()
    state = {
        key: torch.empty(value.size, dtype=value.properties.dtype)
        for key, value in metadata.state_dict_metadata.items()
        if key.startswith(prefix) and hasattr(value, "size")
    }
    dcp.load(
        state,
        storage_reader=reader,
        planner=DefaultLoadPlanner(allow_partial_load=True),
    )
    return {key[len(prefix):]: value for key, value in state.items()}


def sha(sd):
    digest = hashlib.sha256()
    for key in sorted(sd):
        digest.update(key.encode())
        digest.update(sd[key].contiguous().view(torch.uint8).numpy().tobytes())
    return digest.hexdigest()