import simuPOP as sim

import random

class RandomNumOff:
    # a demographic model
    def __init__(self):
        self.numOff = []
    
    def getNumOff(self):
        # return the pre-simulated number of offspring as a generator function
        for item in self.numOff:
            yield item
    
    def __call__(self, pop):
        # define __call__ so that a RandomNumOff object is callable.
        #
        # Each male produce from 1 to 3 offspring. For large population, get the
        # number of males instead of checking the sex of each individual
        self.numOff = [random.randint(1, 3) for ind in pop.individuals() if ind.sex() == sim.MALE]
        # return the total population size
        print('{} mating events with number of offspring {}'.format(len(self.numOff), self.numOff))
        return sum(self.numOff)


pop = sim.Population(10)

# create a demogranic model
numOffModel = RandomNumOff()

pop.evolve(
    preOps=sim.InitSex(),
    matingScheme=sim.RandomMating(
        # the model will be called before mating to deteremine
        # family and population size
        subPopSize=numOffModel,
        # the getNumOff function (generator) returns number of offspring
        # for each mating event
        numOffspring=numOffModel.getNumOff
    ),
    gen=3
)