import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u

from dust_extinction.shapes import P92

fig, ax = plt.subplots()

# generate the curves and plot them
lam = np.logspace(-3.0, 3.0, num=1000)
x = (1.0/lam)/u.micron

ext_model = P92()
ax.plot(1/x,ext_model(x),label='total')

ext_model = P92(FUV_amp=0., NUV_amp=0.0,
                SIL1_amp=0.0, SIL2_amp=0.0, FIR_amp=0.0)
ax.plot(1./x,ext_model(x),label='BKG only')

ext_model = P92(NUV_amp=0.0,
                SIL1_amp=0.0, SIL2_amp=0.0, FIR_amp=0.0)
ax.plot(1./x,ext_model(x),label='BKG+FUV only')

ext_model = P92(FUV_amp=0.,
                SIL1_amp=0.0, SIL2_amp=0.0, FIR_amp=0.0)
ax.plot(1./x,ext_model(x),label='BKG+NUV only')

ext_model = P92(FUV_amp=0., NUV_amp=0.0,
                SIL2_amp=0.0)
ax.plot(1./x,ext_model(x),label='BKG+FIR+SIL1 only')

ext_model = P92(FUV_amp=0., NUV_amp=0.0,
                SIL1_amp=0.0)
ax.plot(1./x,ext_model(x),label='BKG+FIR+SIL2 only')

ext_model = P92(FUV_amp=0., NUV_amp=0.0,
                SIL1_amp=0.0, SIL2_amp=0.0)
ax.plot(1./x,ext_model(x),label='BKG+FIR only')

# Milky Way observed extinction as tabulated by Pei (1992)
MW_x = [0.21, 0.29, 0.45, 0.61, 0.80, 1.11, 1.43, 1.82,
        2.27, 2.50, 2.91, 3.65, 4.00, 4.17, 4.35, 4.57, 4.76,
        5.00, 5.26, 5.56, 5.88, 6.25, 6.71, 7.18, 7.60,
        8.00, 8.50, 9.00, 9.50, 10.00]
MW_x = np.array(MW_x)
MW_exvebv = [-3.02, -2.91, -2.76, -2.58, -2.23, -1.60, -0.78, 0.00,
             1.00, 1.30, 1.80, 3.10, 4.19, 4.90, 5.77, 6.57, 6.23,
             5.52, 4.90, 4.65, 4.60, 4.73, 4.99, 5.36, 5.91,
             6.55, 7.45, 8.45, 9.80, 11.30]
MW_exvebv = np.array(MW_exvebv)
Rv = 3.08
MW_axav = MW_exvebv/Rv + 1.0
ax.plot(1./MW_x, MW_axav, 'o', label='MW Observed')

ax.set_xscale('log')
ax.set_yscale('log')

ax.set_ylim(1e-3,10.)

ax.set_xlabel(r'$\lambda$ [$\mu$m]')
ax.set_ylabel(r'$A(x)/A(V)$')

ax.legend(loc='best')
plt.show()