import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
import astropy.units as u

from dust_extinction.parameter_averages import (CCM89, O94, F99, F04,
                                                VCG04, GCC09, M14, F19, D22,
                                                G23)

fig, ax = plt.subplots(ncols=2, figsize=(10, 4))

# generate the curves and plot them
x = np.arange(1./30., 1./0.0912, 0.001)/u.micron

Rv = 3.1

models = [CCM89, O94, F99, F04, VCG04, GCC09, M14, F19, D22, G23]

for cmodel in models:
   ext_model = cmodel(Rv=Rv)
   indxs, = np.where(np.logical_and(
      x.value >= ext_model.x_range[0],
      x.value <= ext_model.x_range[1]))
   yvals = ext_model(x[indxs])
   ax[0].plot(1./x[indxs], yvals, label=ext_model.__class__.__name__)
   ax[1].plot(1./x[indxs], yvals, label=ext_model.__class__.__name__)

for iax in ax:
   iax.set_xscale('log')
   iax.xaxis.set_major_formatter(ScalarFormatter())

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

ax[0].set_title(f'UV-NIR R(V) = {Rv}')
ax[0].set_xlim(0.08, 3.0)
ax[1].set_title(f'NIR-MIR R(V) = {Rv}')
ax[1].set_xlim(1.0, 32.0)
ax[1].set_ylim(0.0, 0.50)

ax[0].legend(loc='best')
ax[1].legend(loc='best')
plt.tight_layout()
plt.show()