import numpy as np
import scipy.special

# Error function as defined in equation 1
def erf_eq1(x, x0, sx):
    a = 0.5
    b = 0.5
    xx = (x - x0) / np.sqrt(2 * sx**2)
    return a * scipy.special.erf(xx) + b

# Choose the correction appropriate for SDSS apertures
# (if you want to choose for entire galaxies, use rR50 = 2.0)
rR50 = 0.7

coef_aper = {
    3727: [1.130, 0.1302],
    4861: [1.121, 0.1341],
    5007: [1.131, 0.1262],
    6563: [1.121, 0.1340],
    6584: [1.131, 0.1300],
    6725: [1.130, 0.1316],
    }

coef_gal = {
    3727: [1.108, 0.2086],
    4861: [1.078, 0.1944],
    5007: [1.113, 0.1963],
    6563: [1.077, 0.1932],
    6584: [1.098, 0.1944],
    6725: [1.100, 0.1993],
    }

######## This bit just create random measurements to test the code. ########
######## One should replace these by real observed measurements.    ########

# Fix random seed
np.random.seed(42)
    
# Set number of galaxies
Ng = 5

# Create random WHa's, allowing a minimum of 10 and maximum of 20 angstroms.
log_WHa_obs = np.random.uniform(1., np.log10(20.), Ng)

# Create random line fluxes
F_6563_obs = np.random.rand(Ng)
F_6584_obs = np.random.rand(Ng)
######## Done creating random measurements. ########

# Correct line fluxes
F_6563_SFc = F_6563_obs * erf_eq1(log_WHa_obs, *coef_aper[6563])
F_6584_SFc = F_6584_obs * erf_eq1(log_WHa_obs, *coef_aper[6584])

# Check the N2 index, for instance
print('log_WHa_obs =', log_WHa_obs)
print('N2_obs = ', F_6584_obs/F_6563_obs)
print('N2_SFc = ', F_6584_SFc/F_6563_SFc)