import datetime
import math
import numpy as np
from numba import jit, float64

@jit(float64(float64, float64, float64, float64, float64, float64),cache=True,nopython=True)
def sabr_lognormal(alpha, rho, nu, F, K, T):
    A = 1 + (0.25 * (alpha * nu * rho) + nu * nu * (2 - 3 * rho * rho) / 24.) * T
    if F == K:
        VOL = alpha * A
    elif F != K:
        nulogFK = nu * math.log(F/K)
        z = nulogFK / alpha
        x = math.log( ( math.sqrt(1-2*rho*z+z**2) + z - rho ) / (1-rho) )
        VOL = (nulogFK * A) / x
    return VOL

@jit(float64(float64, float64, float64, float64, float64, float64),cache=True,nopython=True)
def sabr_normal(alpha, rho, nu, F, K, T):
    if F == K:
        V = F
        A = 1 + (alpha * alpha / (24. * V * V) + nu * nu * (2 - 3 * rho * rho) / 24.) * T
        VOL = (alpha / V) * A
    elif F != K:
        V = math.sqrt(F * K)
        logFK = math.log(F/K)
        z = (nu/alpha)*V*logFK
        x = math.log( ( math.sqrt(1-2*rho*z+z**2) + z - rho ) / (1-rho) )
        A = 1 + ( (alpha * alpha) / (24. * (V * V)) + ((nu * nu) * (2 - 3 * (rho * rho)) / 24.) ) * T
        logFK2 = logFK * logFK
        B = 1/1920. * logFK2 + 1/24.
        B = 1 + B * logFK2
        VOL = (nu * logFK * A) / (x * B)
    return VOL

@jit(float64(float64, float64, float64, float64, float64, float64, float64),cache=True,nopython=True)
def sabr(alpha, beta, rho, nu, F, K, T):
    if beta == 0.:
        return sabr_normal(alpha, rho, nu, F, K, T)
    elif beta == 1.:
        return sabr_lognormal(alpha, rho, nu, F, K, T)
    else:
        if F == K: # ATM formula
            V = F**(1-beta)
            A = 1 + ( ((1-beta)**2*alpha**2)/(24.*(V**2)) + (alpha*beta*nu*rho) / (4.*V) +
                      ((nu**2)*(2-3*(rho**2))/24.) ) * T
            VOL = (alpha/V)*A
        elif F != K: # not-ATM formula
            V = (F*K)**((1-beta)/2.)
            logFK = math.log(F/K)
            z = (nu/alpha)*V*logFK
            x = math.log( ( math.sqrt(1-2*rho*z+z**2) + z - rho ) / (1-rho) )
            A = 1 + ( ((1-beta)**2*alpha**2)/(24.*(V**2)) + (alpha*beta*nu*rho)/(4.*V) +
                      ((nu**2)*(2-3*(rho**2))/24.) ) * T
            B = 1 + (1/24.)*(((1-beta)*logFK)**2) + (1/1920.)*(((1-beta)*logFK)**4)
            VOL = (nu*logFK*A)/(x*B)
    return VOL

if __name__ == "__main__":
    from analytics.option import BlackSwaption
    from analytics import Index
    from scipy.optimize import least_squares

    underlying = Index.from_name("IG", 28, "5yr")
    underlying.spread = 67.5
    exercise_date = datetime.date(2017, 9, 20)
    option = BlackSwaption(underlying, exercise_date, 70)

    strikes = np.array([50, 55, 57.5, 60, 62.5, 65, 67.5, 70, 75, 80, 85])
    pvs = np.array([44.1, 25.6, 18.9, 14, 10.5, 8.1, 6.4, 5, 3.3, 2.2, 1.5]) * 1e-4

    strikes = np.array([50, 55, 57.5, 60, 62.5, 65, 67.5, 70, 75, 80, 85, 90, 95, 100])
    pvs = np.array([53.65, 37.75, 31.55, 26.45, 22.25, 18.85, 16.15, 13.95, 10.55,
                    8.05, 6.15, 4.65, 3.65, 2.75]) * 1e-4

    def calib(x, option, strikes, pv, beta):
        alpha, rho, nu = x
        F = option.forward_spread
        T = option.T
        r = np.empty_like(strikes)
        for i, K in enumerate(strikes):
            option.strike = K
            option.sigma = sabr(alpha, beta, rho, nu, F, K, T)
            r[i] = option.pv - pv[i]
        return r

    prog = least_squares(calib, (0.3, 0.5, 0.3),
                         bounds=(np.zeros(3), [np.inf, 1, np.inf]),
                         args=(option, strikes, pvs, 1))