import numpy as np
import matplotlib.pyplot as plt
import exo_k as xk

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

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

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

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

kdata = xk.Kdatabase(mols)
corrk_IR = xk.Ktable5d(path=lmdz_corrk, res="17x23", band="IR")
corrk_VI = xk.Ktable5d(path=lmdz_corrk, res="17x23", band="VI")

fig, ax = plt.subplots()
for mol in mols:
    kdata.ktables[mol].plot_spectrum(ax=ax,p=p,t=t,x=x,g=1,label=mol)
corrk_IR.plot_spectrum(ax=ax,p=p,t=t,x=x,g=1,label="IR",color="k",lw=2)
corrk_VI.plot_spectrum(ax=ax,p=p,t=t,x=x,g=1,label="VI",color="k",lw=2,linestyle="dashed")

plt.xlim(0.5,50)
plt.ylim(1e-50,None)
plt.xscale("log")
plt.yscale("log")
plt.legend()
plt.savefig("corrk_bands.png")