import numpy as np

from Prediction.Submatrix import Predictor
from Prediction.VoteResult import SimpleVoteResult


class ComplexSubmatrix(Predictor):

    def separate(self, res: list[SimpleVoteResult]) -> tuple[tuple[list[str], np.ndarray[float]], list[str]]:
        # Values within res are almost always partially observed. That means for example 40% of the votes are counted
        # The percentage of counted votes has to be calculated from self.info and the count within SimpleVoteResult
        # if the percentage is above everage counted percentage, the vote is considered observed completely. The rest is considered unseen
        observed = []
        percentages = []
        not_seen = set(self.Xs.index)

        threshold_for_being_seen = 4/3*np.mean([v.count/self.info.loc[v.identifier].VoteCount for v in res])
        results_with_c = []
        for v in res:
            if (v.count/self.info.loc[v.identifier].VoteCount) >= threshold_for_being_seen:
                results_with_c.append(v)

        for result in results_with_c:
            if result.identifier in self.Xs.index:
                observed.append(result.identifier)
                percentages.append(result.percentage)
                not_seen.remove(result.identifier)

        not_seen = list(set(self.Xs.index).difference(set(observed)))
        return (observed, np.array(percentages)), not_seen