import numpy as np
import matplotlib.pyplot as plt
from scipy.special import gamma

def pNnew_Nold(Nnew, Nold, C, a=0.5):
    """
    P(N|N',a,C) as in Eq. 6 of https://dcc.ligo.org/P2400022
    """
    return gamma(Nnew+Nold+1-a)/(gamma(Nnew+1)*gamma(Nold+1-a))*\
           C**Nnew/(1+C)**(Nnew+Nold+1-a)

# Fix values of C and Nold
C = 1.24  # new value after O4b extension
Nold_BNS = 2
Nold_NSBH = 4

# Plot cumulative probability of a number >=Nnew of detections
# (figure 1 of https://dcc.ligo.org/P2400022)

Nnew = np.arange(30)  # N sufficiently large that p(N)~=0 for large N

pnew_BNS = pNnew_Nold(Nnew, Nold_BNS, C)
cnew_BNS = np.cumsum(pnew_BNS[::-1])[::-1]

pnew_NSBH = pNnew_Nold(Nnew, Nold_NSBH, C)
cnew_NSBH = np.cumsum(pnew_NSBH[::-1])[::-1]

plt.plot(Nnew, cnew_BNS, marker='s', mfc='#FF9C58', mec='r',
         ls='None', label='BNS')
plt.plot(Nnew, cnew_NSBH, marker='o', mfc='#58A1FF', mec='b',
         ls='None', label='NSBH')

plt.xlim(-0.15, 9.15)
plt.ylim(0., 1.1)

plt.xticks(np.arange(10))

plt.xlabel(r'$N_\mathrm{O4b}$')
plt.ylabel(r'Probability that $N\geq N_\mathrm{O4b}$')

plt.legend(loc='upper right')