import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u
from mskpy.util import planck
from grains2 import PlaneParallelIsotropicLTE, ampyroxene50

a = [0.1, 0.33, 1.0, 3.3, 10] * u.um
ap50 = ampyroxene50()
rh = 2.0 * u.au
delta = 1.0 * u.au

lte = PlaneParallelIsotropicLTE(a.value, ap50, rh.value)

i = (lte.wave > 7) * (lte.wave < 14)
wave = lte.wave[i]
Qabs = lte.Qabs[:, i]
F = np.zeros_like(Qabs) * u.Jy
for i in range(len(lte.a)):
    B = planck(lte.T[i], wave, unit="Jy/sr")
    F[i] = (np.pi * a[i]**2 * Qabs[i] * B * u.sr / delta**2).to("Jy")

fig, ax = plt.subplots()
for i in range(len(F)):
    ax.plot(wave, F[i], label=a[i])
ax.legend()
plt.setp(ax, xlabel="Wavelength (μm)", ylabel=r"$F_\nu$ (Jy)", yscale="log")