'''
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
'''
import math
import numpy as np
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 lookup(lookup_value,lookup_vector,result_vector): # find result for the lookup value
   for i, test in enumerate(lookup_vector):
      if i == test:
         return result_vector[i]
   return "NA"

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

A = 1/12
B = A + 1/6
C = A + 2/6
D = A + 3/6
E = A + 4/6
F = A + 5/6
scores = [A,B,C,D,E,F]
RIFvalues = scores # Use the same table for scores and for the midtpoints
characters = ["A","B","C","D","E","F"]
low = 0.0025
medium = 0.01
high = 0.04
def getScore(ch):
   global scores,characters
   return lookup(ch,characters,scores)

Score1_1 = A         # Score of level 1 RIF #1
Score1_2 = D         # Score of level 1 RIF #1
Score2	 = A         # Score of level 2 RIF 
w_1	 = 0.2       # Weight of level 1 RIF #1
w_2	 = 0.8       # Weight of level 1 RIF #2
V_S1_1	 = medium    # Variance of score of level 1 RIF #1 
V_S1_2	 = high      # Variance of score of level 1 RIF #2 
V_S2	 = high      # Variance of score of level 2 RIF  
V_P1	 = low       # Variance of influence level 2 RIF has on level 1 RIF #1 
V_P2	 = low       # Variance of influence level 2 RIF has on level 1 RIF #2 

[alphaP,betaP] = posterior(Score2,V_S2,0.5,0.5) # Parent posterior, using Jefferys prior
parentPosterior = np.zeros(6)
for i in range(6):
   parentPosterior[i] = betaDiscrete(i,alphaP,betaP)

print("\nDiscretized parent 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, i.e, A, B, C, etc.
   # Prior for children, i.e., RIFs on level 1, given parent = p
   beta0C1 = (p * (1 - p) / V_P1 - 1) * (1 - p)
   alpha0C1 = p * beta0C1 / (1 - p)
   beta0C2 = (p * (1 - p) / V_P2 - 1) * (1 - p)
   alpha0C2 = p * beta0C2 / (1 - p)
   # Posterior for children, i.e., RIFs on level 1
   [alphaCh1,betaCh1] = posterior(Score1_1,V_S1_1,alpha0C1,beta0C1)
   [alphaCh2,betaCh2] = posterior(Score1_2,V_S1_2,alpha0C2,beta0C2)   
   for ch1 in range(6):
      for ch2 in range(6):
         mu += (w_1*RIFvalues[ch1]+w_2*RIFvalues[ch2]) *\
                 betaDiscrete(ch1,alphaCh1,betaCh1)* \
                    betaDiscrete(ch2,alphaCh2,betaCh2) * \
                       parentPosterior[i]         
         sm += betaDiscrete(ch1,alphaCh1,betaCh1)* \
                    betaDiscrete(ch2,alphaCh2,betaCh2) * \
                       parentPosterior[i]               
# print(mu,sm)
sig2 = 0
# Calculate variance:
for i, p in enumerate(RIFvalues): # all possible values for the parent
   beta0C1 = (p * (1 - p) / V_P1 - 1) * (1 - p)
   alpha0C1 = p * beta0C1 / (1 - p)
   beta0C2 = (p * (1 - p) / V_P2 - 1) * (1 - p)
   alpha0C2 = p * beta0C2 / (1 - p)
   # Posterior for both children, given parent = p
   [alphaCh1,betaCh1] = posterior(Score1_1,V_S1_1,alpha0C1,beta0C1)
   [alphaCh2,betaCh2] = posterior(Score1_2,V_S1_2,alpha0C2,beta0C2)    
   for ch1 in range(6):
      for ch2 in range(6):
         sig2 += ((w_1*RIFvalues[ch1]+w_2*RIFvalues[ch2])-mu)**2 *\
                 betaDiscrete(ch1,alphaCh1,betaCh1)* \
                    betaDiscrete(ch2,alphaCh2,betaCh2) * \
                       parentPosterior[i]
         
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] + ":","{:.2%}".format(betaDiscrete(i,alpha,beta)))