#%%
import numpy as np
from scipy.integrate import quad
from scipy.special import i0, i0e, erf

import matplotlib.pyplot as plt

d_TA = 78.125                                                 # 1 TA = 78.125 m (see the standard)
K = 1283
d_k  = d_TA * np.arange(0,K)                                  # d_k = k*d_TA with k \in {0,1,...,1282}
e_k  = np.concatenate(([0], (d_k[1:]+d_k[:-1])/2, [np.inf]))  # Distance edges

def p_sk_exact(kk):
    integrand = lambda d, d_hat, σ, λ, e: d/σ**2 * np.exp(-(d-d_hat)**2/(2*σ**2)) * i0e(d*d_hat/σ**2) * (1-np.exp(-λ*(e-d)))

    return np.vectorize(
        lambda d_hat, σ, λ: quad(integrand, 0, e_k[kk+1], args=(d_hat, σ, λ, e_k[kk+1]))[0] - quad(integrand, 0, e_k[kk], args=(d_hat, σ, λ, e_k[kk]))[0]
    )

def p_sk_apprx(kk):
    term = np.vectorize(lambda d_hat, σ, λ, x, ee: \
        0.5*(  erf((x-d_hat)/(np.sqrt(2)*σ)) + np.exp(d_hat*λ+0.5*λ**2*σ**2-ee*λ)*erf((d_hat+λ*σ**2-x)/(np.sqrt(2)*σ))  ) \
    )

    return lambda d_hat, σ, λ, e_kp1=e_k[kk+1], e_k=e_k[kk]: \
        term(d_hat, σ, λ, x=e_kp1, ee=e_kp1) - term(d_hat, σ, λ, x=0, ee=e_kp1) \
        -(term(d_hat, σ, λ, x=e_k, ee=e_k) - term(d_hat, σ, λ, x=0, ee=e_k))

def p_sk_exact_no_mp(kk):
    integrand = lambda d, d_hat, σ, e: d/σ**2 * np.exp(-(d-d_hat)**2/(2*σ**2)) * i0e(d*d_hat/σ**2)

    return np.vectorize(
        lambda d_hat, σ: quad(integrand, 0, e_k[kk+1], args=(d_hat, σ, e_k[kk+1]))[0] - quad(integrand, 0, e_k[kk], args=(d_hat, σ, e_k[kk]))[0]
    )

def p_sk_apprx_no_mp(kk):
    term = np.vectorize(lambda d_hat, σ, x, ee: \
        0.5*(  erf((x-d_hat)/(np.sqrt(2)*σ)) ) \
    )

    return lambda d_hat, σ, e_kp1=e_k[kk+1], e_k=e_k[kk]: \
        term(d_hat, σ, x=e_kp1, ee=e_kp1) - term(d_hat, σ, x=0, ee=e_kp1) \
        -(term(d_hat, σ, x=e_k, ee=e_k) - term(d_hat, σ, x=0, ee=e_k))