from pathlib import Path
import torch

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    StateDictType,
    FullStateDictConfig,  # general model non-sharded, non-flattened params
)


from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


# create singleton saving policies to avoid making over and over
fullstate_save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)


def save_model_checkpoint(filepath, model, rank, verbose=True):
    """saving model via rank0 cpu streaming and full_state_dict"""

    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, fullstate_save_policy):
        cpu_state = model.state_dict()

    if verbose:
        print(f"saving process: rank {rank}  done w model state_dict\n")

    if rank == 0:
        print(f"--> saving model  to {filepath}...")
        # save model
        torch.save(cpu_state, filepath)


def load_model_checkpoint(filepath, model, rank, verbose=True):
    """load local checkpoint to rank0 cpu
    must be called * before * passing to FSDP"""

    if rank != 0:
        return

    model_checkpoint = torch.load(filepath)
    # integrate into loaded model
    model.load_state_dict(model_checkpoint)

    if verbose:
        print("model checkpoint loaded to rank0 cpu")


def save_optimizer_checkpoint(filepath, model, optimizer, rank, verbose=True):
    """save optimizer state via full state dict"""

    if verbose:
        print(f"--> optim state call on rank {rank}\n")

    # pull all sharded optimizer states to rank0 cpu...

    optim_state = FSDP.full_optim_state_dict(model, optimizer)

    if verbose:
        print(f"optim state dict ready on {rank} and len of {len(optim_state)}\n")

    if rank == 0:
        print("--> saving optimizer state...")
        torch.save(optim_state, filepath)
        print(f"--> saved {filepath} to disk")


def load_optimizer_checkpoint(filepath, model, optimizer, rank, cfg):
    """load an fdsp optimizer full_state checkpoint using scatter method
    this ensures only rank 0 loads the optimizer state dict and scatters to other ranks
    """

    opt_file_path = Path.cwd() / cfg.checkpoint_folder / cfg.optimizer_checkpoint_file
    full_osd = None

    if rank == 0:
        full_osd = torch.load(filepath)

        if cfg.verbose:
            print("loaded full osd on rank 0")

    # called from all ranks, though only rank0 has a valid param for full_osd
    # FIXME, might need to be `optimizer` instead of `model`
    sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model, optim=optimizer)

    if cfg.verbose:
        print(f"optimizer shard loaded on rank {rank}")