'''
This code is developed by professor Jørn Vatn. The code can be used
by students and employees at NTNU. The code comes with no warranties.

You are not allowed to use the code for commercial purposes without
written permission from Jørn Vatn
'''
# Compared to ex3_5, the code has been made more efficient on various aspects
import math
import numpy as np
import combLib as c   # Library to find all combinations
def betaDiscrete(i,a,b):  # Get point probability, i = 0:5, discretizing beta(a,b)
   from scipy.stats import beta
   return beta.cdf((i+1)/6,a,b)-beta.cdf(i/6,a,b)

def priorLevel1(p, V_P): # Prior of level 1 RIFs, p = parent value, V_P = variance of parent influence
   beta0 = (p * (1 - p) / V_P - 1) * (1 - p)
   alpha0 = p * beta0 / (1 - p)
   return [alpha0,beta0]

def posterior(s,Vs,prior): # Posterior distribution, socre = s, variance of score = Vs
   [alpha0,beta0] = prior
   alpha = alpha0+s**2*(1-s)/Vs 
   beta = beta0+s*(1-s)**2/Vs    
   return [alpha,beta]

Jeffrey = [0.5,0.5]
A = 1/12
B = A + 1/6
C = A + 2/6
D = A + 3/6
E = A + 4/6
F = A + 5/6
RIFvalues = [A,B,C,D,E,F]
characters = ["A","B","C","D","E","F"]
low = 0.0025
medium = 0.01
high = 0.04

Score1 = [A,D]          # Score of level 1 RIF #1, #2, etc.
Score2 = A              # Score of level 2 RIF 
w      = [0.2,0.8]      # Weights of level 1 RIF #1, #2, etc.
V_S1   = [medium,high]  # Variances of scores of level 1 RIF #1, #2, etc. 
V_S2   = high           # Variance of score of level 2 RIF
V_P    = [low,low]      # Variance of influence level 2 RIF has on level 1 RIF #1, #2, etc. 
nChildren = len(Score1) # Number of children
[alphaP,betaP] = posterior(Score2,V_S2,Jeffrey) # Parent posterior, using Jefferys prior
parentPosterior = np.zeros(6)
for i in range(6):
   parentPosterior[i] = betaDiscrete(i,alphaP,betaP)

print("\nParent posterior:",parentPosterior)

#Calculate unconditional mean for the weighted sum
mu = 0
sm = 0 # Just to verify that probabilities sum to 1
for i, p in enumerate(RIFvalues): # all possible values for the parent
   # Note: Children are only stochastically indpenendent for a given value of the parent
   # Vector of posteriors for the children, given parent = p
   posteriorChildren = []
   for j, s in enumerate(Score1):
      # Prior for child j:
      prior = priorLevel1(p,V_P[j])
      # Add posterior for child j to vector of posteriors for children:
      posteriorChildren.append(posterior(s,V_S1[j],prior))      
   combinations = c.initComb(nChildren) # Initialize vector of combinations
   while c.nxtComb(combinations): # Run through all combinations    
      prob = parentPosterior[i]
      m_j = 0
      for j, w_j in enumerate(w):
         m_j += w_j*RIFvalues[combinations[j]]         
         [alpha,beta] = posteriorChildren[j]
         prob *= betaDiscrete(combinations[j],alpha,beta)      
      mu += m_j * prob
sig2 = 0
# Calculate variance:
for i, p in enumerate(RIFvalues): # all possible values for the parent 
   # Posterior for both children, given parent = p
   posteriorChildren = []
   for j, s in enumerate(Score1):
      # Prior for child j:
      prior = priorLevel1(p,V_P[j])
      # Add posterior for child j to list of posteriors for children:
      posteriorChildren.append(posterior(s,V_S1[j],prior))      
   combinations = c.initComb(nChildren) # Initialize vector of combinations
   while c.nxtComb(combinations): # Run through all combinations    
      prob = parentPosterior[i]
      m_j = 0
      for j, w_j in enumerate(w):
         m_j += w_j*RIFvalues[combinations[j]]         
         [alpha,beta] = posteriorChildren[j]
         prob *= betaDiscrete(combinations[j],alpha,beta)  
      sig2 += (m_j-mu)**2 * prob      
sigma = math.sqrt(sig2)
print("\nUnconditional mean and SD for weighted sum:")
print("Expected value =",round(mu,3))
print("Standard deviation =",round(sigma,3),"\n")      
# Calculate beta-distribution parameters given mean (mu) and standard devation (sigma)
alpha =((1 - mu) / sigma**2 - 1 / mu) * mu ** 2
beta = alpha * (1 / mu - 1)
print("\nUnconditional point distribution weighted sum:")
for i in range(6):
   print(characters[i] +"("+str(round(RIFvalues[i],2)) + "):","{:.2%}".format(betaDiscrete(i,alpha,beta)))