"""Base functions and classes used through the builder package."""

from typing import Tuple, Union, Any


# Shortcuts for typing
rn2_type = Tuple[int, int]
rn3_type = Tuple[int, int, int]
rn_type = Union[rn2_type, rn3_type]
pos2_type = Tuple[float, float]
pos3_type = Tuple[float, float, float]
pos_type = Union[pos2_type, pos3_type]
pbc2_type = Tuple[bool, bool]
pbc3_type = Tuple[bool, bool, bool]
pbc_type = Union[pbc2_type, pbc3_type]
pair_type = Tuple[int, int]
id_pc_type = Tuple[int, int, int, int]
hop_term_type = Tuple[int, int, int, int, int, complex]


def verify_rn(rn: rn_type, complete_item: int = 0) -> rn3_type:
    """
    Check and complete cell index.
    :param rn: incoming cell index
    :param complete_item: item to append to rn if its length is 2
    :return: completed cell index
    """
    rn = tuple(rn)
    len_rn = len(rn)
    assert len_rn == 2 or len_rn == 3, f"Length of rn {rn} should be 2 or 3"
    if len_rn == 2:
        return rn[0], rn[1], complete_item
    else:
        return rn


def verify_pos(pos: pos_type, complete_item: float = 0.0) -> pos3_type:
    """
    Check and complete orbital position.
    :param pos: incoming orbital position
    :param complete_item: item to append to pos if its length is 2
    :return: completed orbital position
    """
    pos = tuple(pos)
    len_pos = len(pos)
    assert len_pos == 2 or len_pos == 3, f"Length of pos {pos} should be 2 or 3"
    if len_pos == 2:
        return pos[0], pos[1], complete_item
    else:
        return pos


def verify_pbc(pbc: pbc_type, complete_item: bool = False) -> pbc3_type:
    """
    Check and complete periodic boundary condition.
    :param pbc: incoming periodic condition
    :param complete_item: item to append to pbc if its length is 2
    :return: completed periodic boundary condition
    """
    pbc = tuple(pbc)
    len_pbc = len(pbc)
    assert len_pbc == 2 or len_pbc == 3, f"Length of pbc {pbc} should be 2 or 3"
    if len_pbc == 2:
        return pbc[0], pbc[1], complete_item
    else:
        return pbc


class Lockable:
    """
    Base class for all lockable objects.

    NOTE: DO NOT use set for __subscribers since it requires subscribers to
    implement the __hash__ method.

    Attributes
    ----------
    __free: bool
        whether the object is unlocked
    __subscribers: List[Any]
        subscribers of this object
    """
    def __init__(self) -> None:
        self.__free = True
        self.__subscribers = []

    def lock(self) -> None:
        """
        Lock the object. Modifications are not allowed then unless the 'unlock'
        method is called.
        :return: None
        """
        self.__free = False

    def unlock(self) -> None:
        """
        Unlock the object. Modifications are allowed then.
        :return: None
        """
        self.__free = True

    def check_editable(self) -> None:
        """
        Check if the object is unlocked for modification.
        :return: None
        """
        assert self.__free, "Modifying a locked object"

    def add_subscriber(self, sub: Any) -> None:
        """
        Add a subscriber
        :param sub: the new subscriber
        :return: None
        """
        self.__subscribers.append(sub)

    def update(self) -> None:
        """
        Call subscribers to update their data.
        :return: None
        """
        for sub in self.__subscribers:
            # If the subscriber has its own subscribers, call them recursively.
            if isinstance(sub, Lockable):
                sub.update()
            # End of recursive subscribing chain at Sample level.
            elif hasattr(sub, "init_array"):
                sub.init_array(force_init=True)
