import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

# --- VOI calculation function ---
def voi_calc_GaussProj(mu, rho, tau, H=None):
    if H is None:
        H = np.eye(2)

    Sigma = np.array([[1, rho], [rho, 1]])
    R = tau**2 * np.eye(H.shape[0])
    S = Sigma @ H.T @ np.linalg.inv(H @ Sigma @ H.T + R) @ H @ Sigma

    mf=np.sum(mu)
    sf=np.sqrt(np.sum(S))
    af=mf/sf
    VOI = mf * norm.cdf(af) + sf * norm.pdf(af) - np.maximum(mf, 0)
  
    return VOI

# --- Compare VOI of partial (1 place) and total test (both places)
rhoV = np.linspace(0.01, 0.99, 99)

# Play with measurement noise standard deviation
tau = 0.3
#tau=1
#tau=0.1

mu = np.zeros(2)
H = np.array([[1, 0]])

VOI_1 = np.zeros_like(rhoV)
VOI_p1 = np.zeros_like(rhoV)

for i, rho in enumerate(rhoV):
    VOI_1[i] = voi_calc_GaussProj(mu, rho, tau)
    VOI_p1[i] = voi_calc_GaussProj(mu, rho, tau, H)

# --- Plot VOI vs correlation ---
plt.figure(1)
plt.plot(rhoV, VOI_1, 'k.', label='Total test')
plt.plot(rhoV, VOI_p1, 'b--', label='Partial test')
plt.legend(loc='lower right')
plt.xlabel('Correlation parameter')
plt.ylabel('VOI')
plt.title('VOI vs Correlation')
plt.grid(True)

# --- Decision Regions ---
rho = 0.7
VOI2 = voi_calc_GaussProj(mu, rho, tau)
VOI1 = voi_calc_GaussProj(mu, rho, tau, H)

# Price region plot for this set correlation rho
plt.figure(2)
plt.plot([VOI1, VOI1], [VOI2-VOI1, 2*(VOI2-VOI1)], 'k')
plt.plot([0, VOI1], [VOI2-VOI1, VOI2-VOI1], 'k')
plt.plot([VOI1, VOI2], [VOI2-VOI1, 0], 'k')
plt.text(VOI1/3, (VOI2-VOI1)/3, 'Total')
plt.text(VOI1/2, 1.1*(VOI2-VOI1), 'Partial')
plt.text(VOI1*1.1, (VOI2-VOI1), 'No testing')
plt.xlabel('Price of first test')
plt.ylabel('Price of second test')
plt.title('Decision Region Plot 1')