from dataclasses import dataclass
from typing import Optional, Dict

__all__ = [
    'State',
    'TrainerState'
]


class State(dict):
    r"""
    提供给用户使用的 ``state``，用来记载您的 ``callback`` 实时数据，该 ``state`` 实际上是一个字典，我们通过复用 ``__getattr__`` 方法来实现类似
    类属性的字典调用方式；

    为了实现断点重训，用户应当保证其保存的信息都是可序列化的；

    推荐的使用方式::

        >>> state = State()
        >>> state["best_accuracy"] = 0.9
        >>> print(state["best_accuracy"])
        or
        >>> print(state.best_accuracy)
    """

    __slots__ = ()  # 用户不应当使用 state.name = "name" 来使用此类，因此我们限制用户不可自己对该类设置属性，但是可以通过属性访问字典；

    def __init__(self, *args, **kwargs):
        super(State, self).__init__(*args, **kwargs)

    def __getattr__(self, item):
        if item in self:
            _value = self[item]
            if isinstance(_value, dict):
                return State(_value)
            else:
                return _value
        else:
            raise ValueError(f"key '{item}' is not existed!")

@dataclass
class TrainerState:
    r"""
    该类用于我们 fastNLP 自己内部为了训练流程所记录的一些状态，当然是要暴露给用户给用户使用的；
    我们保存的 state 大部分上是 trainer 断点重训 需要重新加载的；
    专属于 `Trainer` 的状态记载的类；

    :param n_epochs: 训练过程中总共的 epoch 的数量；
    :param cur_epoch_idx: 当前正在运行第几个 epoch；
    :param global_forward_batches: 当前模型总共 forward 了多少个 step；
    :param batch_idx_in_epoch: 训练中在当前 epoch 的第几个 step；
    :param num_batches_per_epoch: 每一个 epoch 会 forward 多少个 step；
    :param n_batches: 完整训练过程会 forward 的 step 数量，注意 ``n_batches = num_batches_per_epoch * n_epochs`` ；
    """
    n_epochs: Optional[int] = None  # 无论如何重新算

    cur_epoch_idx: Optional[int] = 0  # 断点重训; 仅当 resume=False 时为0；
    global_forward_batches: Optional[int] = 0  # 断点重训

    batch_idx_in_epoch: Optional[int] = 0  # 断点重训

    num_batches_per_epoch: Optional[int] = None  # 无论如何重新算

    n_batches: Optional[int] = None  # 无论如何重新算

    def state_dict(self) -> Dict:
        r"""
        :return: 用于断点重训来保存的状态字典；
        """
        return {"cur_epoch_idx": self.cur_epoch_idx, "global_forward_batches": self.global_forward_batches,
                "batch_idx_in_epoch": self.batch_idx_in_epoch}

    def load_state_dict(self, state_dict: Dict):
        r"""
        用于断点重训来重新加载保存的状态字典；

        :param state_dict: 用于加载的状态字典；
        """
        for key in state_dict:
            assert key in {"cur_epoch_idx", "global_forward_batches", "batch_idx_in_epoch"}, "Wrong state_dict for `TrainerState`."
            setattr(self, key, state_dict[key])


