from math import log, sqrt, erf
from numba import jit, float64, boolean
from scipy.stats import norm
import math


def d1(F, K, sigma, T):
    return (log(F / K) + sigma**2 * T / 2) / (sigma * math.sqrt(T))


def d2(F, K, sigma, T):
    return d1(F, K, sigma, T) - sigma * math.sqrt(T)


@jit(cache=True, nopython=True)
def d12(F, K, sigma, T):
    sigmaT = sigma * sqrt(T)
    d1 = log(F / K) / sigmaT
    d2 = d1
    d1 += 0.5 * sigmaT
    d2 -= 0.5 * sigmaT
    return d1, d2


@jit(float64(float64), cache=True, nopython=True)
def cnd_erf(d):
    """ 2 * Phi where Phi is the cdf of a Normal """
    RSQRT2 = 0.7071067811865475
    return 1 + erf(RSQRT2 * d)


@jit(float64(float64, float64, float64, float64, boolean), cache=True, nopython=True)
def black(F, K, T, sigma, payer=True):
    d1, d2 = d12(F, K, sigma, T)
    if payer:
        return 0.5 * (F * cnd_erf(d1) - K * cnd_erf(d2))
    else:
        return 0.5 * (K * cnd_erf(-d2) - F * cnd_erf(-d1))


@jit(float64(float64, float64, float64, float64), cache=True, nopython=True)
def Nx(F, K, sigma, T):
    return cnd_erf((log(F/K) - sigma**2 * T / 2) / (sigma * sqrt(T))) / 2


def bachelier(F, K, T, sigma):
    """ Bachelier formula for normal dynamics

    need to multiply by discount factor
    """
    d1 = (F - K) / (sigma * sqrt(T))
    return (0.5 * (F - K) * cnd_erf(d1) + sigma * sqrt(T) * norm.pdf(d1))