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.0) * 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.0 * V * V) + nu * nu * (2 - 3 * rho * rho) / 24.0)
            * 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.0 * (V * V))
                + ((nu * nu) * (2 - 3 * (rho * rho)) / 24.0)
            )
            * T
        )
        logFK2 = logFK * logFK
        B = 1 / 1920.0 * logFK2 + 1 / 24.0
        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.0:
        return sabr_normal(alpha, rho, nu, F, K, T)
    elif beta == 1.0:
        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.0 * (V ** 2))
                    + (alpha * beta * nu * rho) / (4.0 * V)
                    + ((nu ** 2) * (2 - 3 * (rho ** 2)) / 24.0)
                )
                * T
            )
            VOL = (alpha / V) * A
        elif F != K:  # not-ATM formula
            V = (F * K) ** ((1 - beta) / 2.0)
            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.0 * (V ** 2))
                    + (alpha * beta * nu * rho) / (4.0 * V)
                    + ((nu ** 2) * (2 - 3 * (rho ** 2)) / 24.0)
                )
                * T
            )
            B = (
                1
                + (1 / 24.0) * (((1 - beta) * logFK) ** 2)
                + (1 / 1920.0) * (((1 - beta) * logFK) ** 4)
            )
            VOL = (nu * logFK * A) / (x * B)
    return VOL


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

    underlying = CreditIndex("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),
    )