# cython: language_level=3, cdivision=True, boundscheck=False, wraparound=False
from libc.math cimport log, exp, sqrt
from .black cimport cnd_erf
cimport cython
import numpy as np
cimport numpy as np

cdef api double h1(int n, double* data) nogil:
    # z = (y - mu_y) / sigma_y
    cdef:
       double z = data[0]
       double K = data[1]
       double S1 = data[2]
       double S2 = data[3]
       double mu_x = data[4]
       double mu_y = data[5]
       double sigma_x = data[6]
       double sigma_y = data[7]
       double rho = data[8]
       double u1, u2, Ktilde, v, v2, x

    u1 = mu_x + rho * sigma_x * z
    Ktilde = K + S2 * exp(mu_y + sigma_y * z)
    u2 = log(Ktilde / S1)

    v = sigma_x * sqrt(1 - rho * rho)
    v2 = sigma_x * sigma_x * (1 - rho * rho)
    x = (u1 - u2) / v
    return (
        0.5
        * (S1 * exp(u1 + 0.5 * v2) * cnd_erf(x + v) - Ktilde * cnd_erf(x))
        * exp(-0.5 * z * z)
    )


cpdef np.ndarray h_call(double[::1] z, double K, double S1, double S2, double mu_x, double mu_y, double sigma_x, double sigma_y, double rho):
    # conditionned on S2, integral wrt S1
    # z = (y - mu_y) / sigma_y
    cdef:
       np.ndarray[np.float_t, ndim=1] r = np.empty_like(z)
       double u1, u2, Ktilde, x
       double v = sigma_x * sqrt(1 - rho * rho)
       double v2 = sigma_x * sigma_x * (1 - rho * rho)
       int i

    for i in range(z.shape[0]):
        u1 = mu_x + rho * sigma_x * z[i]
        Ktilde = K + S2 * exp(mu_y + sigma_y * z[i])
        u2 = log(Ktilde / S1)
        x = (u1 - u2) / v
        r[i] = 0.5 * (S1 * exp(u1 + 0.5 * v2) * cnd_erf(x + v) - Ktilde * cnd_erf(x))
    return r

cpdef double[::1] h_put(double[::1] z, double K, double S1, double S2, double mu_x, double mu_y, double sigma_x, double sigma_y, double rho):
    # z = (y - mu_y) / sigma_y
    cdef:
       double[::1] r = np.empty_like(z)
       double u1, u2, Ktilde, x
       double v = sigma_x * sqrt(1 - rho * rho)
       double v2 = sigma_x * sigma_x * (1 - rho * rho)
       int i

    for i in range(z.shape[0]):
        u1 = mu_x + rho * sigma_x * z[i]
        Ktilde = K + S2 * exp(mu_y + sigma_y * z[i])
        u2 = log(Ktilde / S1)
        x = (u2 - u1) / v
        r[i] = 0.5 * (Ktilde * cnd_erf(x) - S1 * exp(u1 + 0.5 * v2) * cnd_erf(x - v))
    return r