from collections import Counter
from helpers import Parents

assert round_values(parent_genotypes(
    parents_phenotypes=Parents( # parents are specified pheontypically
        ('rough', 'long haired', 'black'),
        ('rough', 'short haired', 'white'),
    ),
    offspring_phenotypes=list(Counter({ # offspring as well
        ('rough', 'short haired', 'black'): 15,
        ('rough', 'long haired', 'black'): 13,
        ('smooth', 'short haired', 'black'): 4,
        ('smooth', 'long haired', 'black'): 4,
    }).elements()),
    structure=( # show each of the traits, dominant first. genotype is created with first lettter of word
        ('rough', 'smooth'),
        ('short haired', 'long haired'),
        ('black', 'white'),
    ),
), 1) == Counter({'RrssBB (x) RrSsbb': 1.0, 'RrssBb (x) RrSsbb': 0.0}) # map genotypes of parents to probability

from helpers import *

from itertools import repeat
from collections import Counter

def parent_genotypes(parents_phenotypes, offspring_phenotypes, structure):
    '''
    Returns the probabiblity of different parent genotypes, based on the phenotypes of the known children
    and parents.
    '''

    parents_genotypes_probability = Counter()

    # lets try a brute force technique. get a combination of every possible parents genotypes and
    # try them on the parents and offspring phenotypes to see if they match.
    for trial_parents_genotypes in all_genotype_parents_from_structure(structure):
        # first make sure that the parents work
        trial_parents_phenotypes = map(phenotype_from_genotype, trial_parents_genotypes.elements(), repeat(structure))
        if Counter(trial_parents_phenotypes) == Counter(parents_phenotypes):
            # then we want to see if the real offspring are a subset of the children of those parents
            trial_offspring_genotypes = offspring_genotypes(trial_parents_genotypes.elements())
            trial_offspring_phenotypes = Counter(map(phenotype_from_genotype, trial_offspring_genotypes.elements(), repeat(structure)))
            if set(offspring_phenotypes).issubset(set(trial_offspring_phenotypes)):
                # then we know that those parents could actually be the right ones, however they might not be,
                # so we also need to compute, how probable it is.

                # one way to look at it is, if the `trial_parents` had `len(offspring_phenotypes)` kids, what percentage
                # of the time would the phenotypes of their kids be `offspring_phenotypes`?

                # that gives us the probability of getting the actual litter we have with these parents.

                # then we can compare that probability to the other options, and see which one is more likely.
                chance_that_offspring_came_from_trial_parents = 1

                trial_offspring_phenotypes_percentiles = normalize_values(trial_offspring_phenotypes)
                for offspring_phenotype in offspring_phenotypes:
                    # calculate the chance of getting each child, and since we are talking about succesion we multiply them
                    chance_that_offspring_came_from_trial_parents *= trial_offspring_phenotypes_percentiles[offspring_phenotype]

                parents_genotypes_probability[pprint_parents_genotypes(trial_parents_genotypes.elements())] = chance_that_offspring_came_from_trial_parents

    return normalize_values(parents_genotypes_probability)

def n_dom_child(n):
    return parent_genotypes(
        parents_phenotypes=Parents(
            ('dom',),
            ('rec',),
        ),
        offspring_phenotypes=list(Counter({('dom',): n}).elements()),
        structure=(
            ('dom', 'rec'),
        ),
    )

n_dom_child(1)

[[i, n_dom_child(i).most_common()] for i in range(10)]