import numpy as np

from numpy       import pi, sqrt, log, exp
from scipy.stats import norm

#
# Normal cumulative distribution function, with extension
# for complex argument with small imaginary component
#

def norm_cdf(x):
    if not isinstance(x, np.ndarray):
        xr = x.real
        xi = x.imag
        if abs(xi) > 1.0e-10:
            raise ValueError('imag(x) too large in norm_cdf(x)')

        ncf = norm.cdf(xr)
        if abs(xi) > 0:
            ncf = ncf + 1.0j*xi*norm.pdf(xr)
    else:
        xr = np.real(x)
        xi = np.imag(x)
        if any(abs(xi) > 1.0e-10):
            raise ValueError('imag(x) too large in norm_cdf(x)')

        ncf = norm.cdf(xr)
        if any(abs(xi) > 0):
            ncf = ncf + 1.0j*xi*norm.pdf(xr)

    return ncf

#
# function V = digital_put(r,sigma,T,S,K,opt)
#
# Black-Scholes digital put option solution
#
# r     - interest rate
# sigma - volatility
# T     - time interval
# S     - asset value(s)  (float or flattened numpy array)
# K     - strike price(s) (float or flattened numpy array)
# opt   - 'value', 'delta', 'gamma' or 'vega'
# V     - option value(s) (float or flattened numpy array)
#

def digital_put(r,sigma,T,S,K,opt):

    S  = S + 1e-100    # avoids problems with S=0
    K  = K + 1e-100    # avoids problems with K=0

    d2 = ( log(S) - log(K) + (r-0.5*sigma**2)*T ) / (sigma*sqrt(T))

    if opt == 'value':
        V = exp(-r*T)*(1-norm_cdf(d2))

    elif opt == 'delta':
        V = - exp(-r*T)*(exp(-0.5*d2**2)/sqrt(2*pi)) / (sigma*sqrt(T)*S)

    elif opt == 'gamma':
        V = - exp(-r*T)*(exp(-0.5*d2**2)/sqrt(2*pi)) \
             *(-d2/(sigma*sqrt(T)*S) - 1/S)/(sigma*sqrt(T)*S)
 
    elif opt == 'vega':
        V = - exp(-r*T)*(exp(-0.5*d2**2)/sqrt(2*pi))*(-d2/sigma-sqrt(T))

    else:
        raise ValueError('invalid value for opt -- must be "value", "delta", "gamma", "vega"')

    return V