/* 

  Rosalind problem Calculating Expected Offspring in Picat.

  https://rosalind.info/problems/iev/
  """
  Problem
  For a random variable X taking integer values between 1 and n, the expected value of X is 
  E(X)=∑nk=1k×Pr(X=k). The expected value offers us a way of taking the long-term average of a 
  random variable over a large number of trials.

  As a motivating example, let X be the number on a six-sided die. Over a large number of rolls, 
  we should expect to  obtain an average of 3.5 on the die (even though it's not possible to roll a 3.5). 
  The formula for expected value confirms that E(X)=∑6k=1k×Pr(X=k)=3.5.

  More generally, a random variable for which every one of a number of equally spaced outcomes has the 
  same probability is called a uniform random variable (in the die example, this "equal spacing" 
  is equal to 1). We can generalize our die example to find that if X is a uniform random variable with 
  minimum possible value a and maximum possible value b, then E(X)=a+b2. You may also wish to verify that 
  for the dice example, if Y is the random variable associated with the outcome of a second die roll, then E(X+Y)=7.

  Given: Six nonnegative integers, each of which does not exceed 20,000. The integers correspond to the number 
  of couples in a population possessing each genotype pairing for a given factor. In order, the six given 
  integers represent the number of couples having the following genotypes:

  1. AA-AA
  2. AA-Aa
  3. AA-aa
  4. Aa-Aa
  5. Aa-aa
  6. aa-aa

  Return: The expected number of offspring displaying the dominant phenotype in the next generation, under 
  the assumption that every couple has exactly two offspring.

  Sample Dataset
     1 0 0 1 0 1
  Sample Output
     3.5
  """

  1. AA-AA:  -> AA, AA, AA, AA  : 4/4
  2. AA-Aa:  -> AA, Aa, AA, Aa  : 4/4
  3. AA-aa:  -> Aa, Aa, Aa, Aa  : 4/4
  4. Aa-Aa:  -> AA, Aa, Aa, aa  : 3/4
  5. Aa-aa:  -> Aa, Aa, aa, aa  : 2/4
  6. aa-aa:  -> aa, aa, aa, aa  : 0/4

  1 of AA-AA  : 2*4/4
  1 of Aa-Aa  : 2*3/4
  1 of aa-aa  : 0
  2+6/4 = 3.5
  

  

  This model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/picat/

*/

import util.

main => go.

go ?=>
  % ["AA-AA","AA-Aa","AA-aa","Aa-Aa","Aa-aa","aa-aa"]]
  % Genotypes = [split(G,'-') : G in findall(G,genotype(_Id,G,_P))],
  % Number of dominant per pair of genotypes
  % DominantP = [4/4,4/4,4/4,3/4,2/4,0/4],
  DominantP = [P : P in findall(P,genotype(_Id,_G,P))],  
  println(DominantP),
  
  L = "1 0 0 1 0 1",
  % L := "16473 18362 17725 18245 17107 17073", % The Challenge problem
  NumOffsprings = 2,
  % println(l=L),  
  X = [I.to_int : I in L.split(" ")],
  println(x=X),  
  
  S = sum([NumOffsprings*X[I]*DominantP[I] : I in 1..6]),
  println(s=S),

  nl.
go => true.

% Shorter
go2 ?=>
  L = "1 0 0 1 0 1",
  println(2*sum(map(*,[1,1,1,0.75,0.5,0],map(to_int,L.split)))),
  
  nl.
go2  => true.

% Calculate the pcts of a dominant offspring
go3 ?=>
  genotype(Id,Pair,_),
  println(Id=Pair),
  [P1,P2] = Pair.split("-"),
  Count = 0,
  foreach(G1 in P1, G2 in P2)
    if G1 == 'A' ; G2 == 'A' then
      Count := Count + 1
    end
  end,
  println(count=Count),
  println(pct=Count/4),
  nl,
  fail,
  nl.
go3 => true.

%
% Encapsulated variant of go3/0
%
go4 ?=>
  println(pcts=pcts()),
  nl.
go4 => true.


% Percentage of getting a dominant offspring per type of genotype pair
%           pair pct
genotype(1,"AA-AA",4/4).
genotype(2,"AA-Aa",4/4).
genotype(3,"AA-aa",4/4).
genotype(4,"Aa-Aa",3/4).
genotype(5,"Aa-aa",2/4).
genotype(6,"aa-aa",0/4).

pcts() = new_map([Pair=pct(Pair) : Pair in findall(Pair,genotype(_Id,Pair,_))]).

pct(Pair) = Pct => 
  T=[[A,B]: {A,B} in zip(Pair.split("-"))],
  Pct = sum([ 1 : P1 in T[1], P2 in T[2], (P1 == 'A' ; P2 == 'A')])/4.