from votekit.pref_profile import RankProfile, ProfileError
from typing import Optional, Callable
from votekit.models import Election
from votekit.elections.election_state import ElectionState
from abc import abstractmethod


class RankingElection(Election[RankProfile]):
    """
    Abstract class for ranking election types. Validates that profile contains ranked ballots.

    Args:
        profile (RankProfile): The initial profile of ballots.
        score_function (Callable[[RankProfile], dict[str, float]], optional):
            A function that converts profiles to a score dictionary mapping candidates to
            their current score. Used in creating ElectionState objects and sorting candidates in
            Round 0. If None, no score dictionary is saved and all candidates are tied in Round 0.
            Defaults to None.
        sort_high_low (bool, optional): How to sort candidates based on `score_function`. True sorts
            from high to low. Defaults to True.

    Attributes:
        election_states (list[ElectionState]): a list of election states, one for each round of
            the election. The list is 0 indexed, so the initial state is stored at index 0, round 1
            at 1, etc.
        score_function (Callable[[RankProfile], dict[str, float]], optional):
            A function that converts profiles to a score dictionary mapping candidates to
            their current score. Used in creating ElectionState objects. Defaults to None.
        length (int): the number of rounds of the election.
    """

    def __init__(
        self,
        profile: RankProfile,
        score_function: Optional[Callable[[RankProfile], dict[str, float]]] = None,
        sort_high_low: bool = True,
    ):
        super().__init__(profile, score_function, sort_high_low)

    def _validate_profile(self, profile: RankProfile):
        """
        Validate that a profile contains appropriate ballots for election. Raises
        TypeError if not.

        Args:
            profile (RankProfile): Profile of ballots.
        """
        if not isinstance(profile, RankProfile):
            raise ProfileError("Profile must be of type RankProfile.")

        assert profile.max_ranking_length is not None
        ranking_cols = [
            f"Ranking_{i}" for i in range(1, profile.max_ranking_length + 1)
        ]
        try:
            np_arr = profile.df[ranking_cols].to_numpy()
            tilde = frozenset("~")
            for row in np_arr:
                if (row == tilde).all():
                    # TODO: Put a better error message here
                    raise RuntimeError
        except (KeyError, RuntimeError):
            for ballot in profile.ballots:
                if ballot.ranking is None:
                    raise TypeError(f"Ballot {ballot} has no ranking.")

    @abstractmethod
    def _run_step(
        self, profile: RankProfile, prev_state: ElectionState, store_states=False
    ) -> RankProfile:
        """
        Run one step of an election from the given profile and previous state.

        Args:
            profile (RankProfile): Profile of ballots.
            prev_state (ElectionState): The previous ElectionState.
            store_states (bool, optional): True if `self.election_states` should be updated with the
                ElectionState generated by this round. This should only be True when used by
                `self._run_election()`. Defaults to False.

        Returns:
            RankProfile: The profile of ballots after the round is completed.
        """
        pass

    @abstractmethod
    def _is_finished(self) -> bool:
        """
        Returns True if election is finished, False if another round is needed.
        """
        pass
