from .tranche_functions import GHquad
from math import exp, sqrt, log
from .blac import cnd_erf

def CMS_spread(T_alpha, X, beta, gamma):
    Z, w = GHquad(100)
    return np.inner(f(Z), w)

def f(v, X, S_alpha_beta, S_alpha_gamma, mu_beta, mu_gamma, T_alpha, rho):
    h = h(v, X, S_alpha_beta, mu_beta, sigma_alpha_beta, T_alpha)
    u = rho * sigma_alpha_gamma * sqrt(T_alpha) * v
    d = sigma_alpha_gamma * sqrt(T_alpha) * sqrt(1 - rho ** 2)
    r = mu_gamma * T_alpha - 0.5 * rho * rho * sigma_alpha_gamma ** 2 * T_alpha + u
    u0 = log(S_alpha_gamma / h) + u
    u1 = u0 + (mu_gamma + (0.5 - rho ** 2) * sigma_alpha_gamma**2) * T_alpha
    u2 = u0 + (mu_gamma - 0.5 * sigma_alpha_gamma**2) * T_alpha
    return 0.5 * (S_alpha_gamma * exp(r) * cnd_erf(u1 / d) - h * cnd_erf(u2 / d))


def h(v, X, S_alpha_beta, mu_beta, sigma_alpha_beta, T_alpha):
    r = (mu_beta - 0.5 * sigma_alpha_beta * sigma_alpha_beta) * T_alpha + \
        sigma_alpha_beta * sqrt(T_alpha) * v
    return X + S_alpha_beta * exp(r)