import math
import numpy
import pytest

from tests.tests_perf.data_generator import *
from tests.tests_perf.test_perf_base import DPNPTestPerfBase


SEED = 7777777
SL, SH = 10.0, 50.0
KL, KH = 10.0, 50.0
TL, TH = 1.0, 2.0
RISK_FREE = 0.1
VOLATILITY = 0.2


def math_erf(x):
    result = numpy.empty(x.shape, dtype=x.dtype)
    for i in range(result.size):
        result[i] = math.erf(x[i])

    return result


# module 'numpy' has no attribute 'erf'
numpy.erf = math_erf


def gen_data(lib, low, high, size):
    return lib.random.uniform(low, high, size)


def black_scholes_put(lib, S, K, T, r, sigma):
    d1 = (lib.log(S / K) + (r + sigma * sigma / 2.) * T) / (sigma * lib.sqrt(T))
    d2 = d1 - sigma * lib.sqrt(T)

    cdf_d1 = (1 + lib.erf(d1 / lib.sqrt(2))) / 2
    cdf_d2 = (1 + lib.erf(d2 / lib.sqrt(2))) / 2

    bs_call = S * cdf_d1 - K * lib.exp(-r * T) * cdf_d2

    return K * lib.exp(-r * T) - S + bs_call


class TestBlackScholes(DPNPTestPerfBase):

    @pytest.mark.parametrize("dtype", [numpy.float64])
    @pytest.mark.parametrize("size", [1024, 2048, 4096, 8192])
    def test_bs_put(self, lib, dtype, size):
        numpy.random.seed(SEED)
        S = gen_data(lib, SL, SH, size)
        K = gen_data(lib, KL, KH, size)
        T = gen_data(lib, TL, TH, size)

        self.dpnp_benchmark("bs_put", lib, dtype, size,
                            lib, S, K, T, RISK_FREE, VOLATILITY,
                            custom_fptr=black_scholes_put)
