import numpy as np
import matplotlib.pyplot as plt
import exo_k as xk
plt.set_cmap("viridis")

### where is the data?
corrk_dir="datadir/corrk_data/"
lmdz_corrk = corrk_dir+"pluton_ch4var_COfix/"
# xk.Settings().set_search_path(corrk_dir+"xsec/taurex_R15000/")

### select pressure, temperature, and mixing ratio to plot
p=1e8
t=300
x=1e-4

mols = [
"C2H4",
"CO",
# "C2H2",
# "CH4",
]

xk.Settings().set_mks(True)
xk.Settings().set_log_interp(False)
xk.Settings().set_case_sensitive(True)

fig, ax = plt.subplots()

xk.Settings().set_search_path("datadir/exomol/taurex_R15000")
kdata = xk.Kdatabase(mols)
kdata.bin_down(xk.wavenumber_grid_R(10000/50,10000/0.5,100))
print(kdata)
for mol in mols:
    kdata.ktables[mol].plot_spectrum(ax=ax,p=p,t=t,x=x,g=1,label=mol+" exomol (original)", linestyle="dashdot")

plt.gca().set_prop_cycle(None)


xk.Settings().set_search_path(corrk_dir+"R500_from_R15000/")
kdata = xk.Kdatabase(mols)
kdata.bin_down(xk.wavenumber_grid_R(10000/50,10000/0.5,100))
print(kdata)
for mol in mols:
    kdata.ktables[mol].plot_spectrum(ax=ax,p=p,t=t,x=x,g=1,label=mol+" exomol (R500)")

isotopologues = ["12C-16O", "12C2-1H4",
                #   "12C-1H4", "12C-1H2"
                ]

plt.gca().set_prop_cycle(None)
xk.Settings().set_search_path(corrk_dir+"R500_from_dace/")
kdata = xk.Kdatabase(isotopologues)
print(kdata)
for mol in mols:
    try:
        kdata.ktables[mol].plot_spectrum(ax=ax,p=p,t=t,x=x,g=1,label=mol+" dace (R500)", dashes=[5,5])
    except:
        print(f"Missing {mol} dace")

plt.xlim(0.5,20)
plt.ylim(1e-40,None)
# plt.xscale("log")
plt.yscale("log")
plt.legend()
plt.savefig(f"corrk_dace_exomol_p{p}_t{t}_x{x}.png")