
import numpy as np
from mapel.core.utils import get_vector


def generate_approval_partylist_votes(num_voters=None, num_candidates=None, g=5):

    num_groups = g

    alphas = get_vector('linear', num_groups)
    sizes = np.random.dirichlet(alphas)
    cumv = np.cumsum(sizes)
    cumv = np.insert(cumv, 0, 0)

    votes = []

    for a in range(1, num_groups+1):
        vote = set()
        for i in range(int(num_candidates*cumv[a-1]), int(num_candidates*cumv[a])):
            vote.add(i)
        for i in range(int(num_candidates * cumv[a - 1]), int(num_candidates * cumv[a])):
            votes.append(vote)
    return votes


def generate_approval_urn_partylist_votes(num_voters=None, num_candidates=None, g=5, alpha=0.1):
    num_groups = g

    party_votes = np.zeros([num_voters])
    urn_size = 1.
    for j in range(num_voters):
        rho = np.random.uniform(0, urn_size)
        if rho <= 1.:
            party_votes[j] = np.random.randint(0, num_groups)
        else:
            party_votes[j] = party_votes[np.random.randint(0, j)]
        urn_size += alpha

    party_size = int(num_candidates/num_groups)
    votes = []

    for i in range(num_voters):
        shift = party_votes[i]*party_size
        vote = set([int(c+shift) for c in range(party_size)])
        votes.append(vote)

    return votes


