#This code generate the KDTree template file that can be directly used.

import numpy as np
import matplotlib.pyplot as plt
import cPickle as pickle
from sklearn.neighbors import KDTree
from scipy.interpolate import splrep, splev
from sedfit.fitter.template import Template

modelDir = "/Users/jinyi/Work/mcmc/Fitter/template/grain_models/crosec/"
silNameList = ["DL84", "amo4", "pyr7"]
graNameList = ["grap"]
silTList = np.arange(len(silNameList)) #The types of silicate dust.
graTList = np.arange(len(graNameList)) #The types of graphite dust.

#"""
#Obtain the silicate template interpolated list
tckSilList = []
parSilList = []
for nsil in silTList:
    tmplDir = modelDir + "{0}.dat".format(silNameList[nsil])
    modelSil = np.loadtxt(tmplDir)
    size_all  = modelSil[:, 0]
    wave_all  = modelSil[:, 1]
    kappa_all = modelSil[:, 2]
    sizeList = np.unique(size_all)
    for sz in sizeList:
        fltr = size_all == sz
        wave  = wave_all[fltr]
        kappa = kappa_all[fltr]
        tck = splrep(wave, kappa)
        tckSilList.append(tck)
        parSilList.append([nsil, sz])
kdtSil = KDTree(parSilList)
print("Silicate dust interpolation finishes!")
modelInfo = {
    "type": silTList, #Silicate dust type
    "size": sizeList,
    "wavelength": wave,
}
parFormat = ["type", "size"]
readMe = '''
This silicate template is generated by prof. Aigen Li's code.
The interpolation is tested well!
'''
silDict = {
    "tckList": tckSilList,
    "kdTree": kdtSil,
    "parList": parSilList,
    "modelInfo": modelInfo,
    "parFormat": parFormat,
    "readMe": readMe
}

#Obtain the graphite template interpolated list
tckGraList = []
parGraList = []
for ngra in graTList:
    tmplDir = modelDir + "{0}.dat".format(graNameList[ngra])
    modelGra = np.loadtxt(tmplDir)
    size_all  = modelGra[:, 0]
    wave_all  = modelGra[:, 1]
    kappa_all = modelGra[:, 2]
    sizeList = np.unique(size_all)
    for sz in sizeList:
        fltr = size_all == sz
        wave  = wave_all[fltr]
        kappa = kappa_all[fltr]
        tck = splrep(wave, kappa)
        tckGraList.append(tck)
        parGraList.append([ngra, sz])
kdtGra = KDTree(parGraList)
print("Graphite dust interpolation finishes!")
modelInfo = {
    "type": graTList, #Silicate dust type
    "size": sizeList,
    "wavelength": wave,
}
parFormat = ["type", "size"]
readMe = '''
This graphite template is generated by prof. Aigen Li's code.
The interpolation is tested well!
'''
graDict = {
    "tckList": tckGraList,
    "kdTree": kdtGra,
    "parList": parGraList,
    "modelInfo": modelInfo,
    "parFormat": parFormat,
    "readMe": readMe
}

#Save the template
templateReadMe = '''
This template file consists two template dicts for silicate and graphite dust
grain. Each of the dicts can be used as input of the Template class in sedfit
package.
'''
dustModel = {
    "Silicate": silDict,
    "Graphite": graDict,
    "readMe": templateReadMe
}
fp = open("/Users/jinyi/Work/mcmc/Fitter/template/dust_xl_kdt.tmplt", "w")
pickle.dump(dustModel, fp)
fp.close()
#"""

##Test the KDTree and the interpolation
fp = open("/Users/jinyi/Work/mcmc/Fitter/template/dust_xl_kdt.tmplt", "r")
dustModel = pickle.load(fp)
fp.close()
#For the astronomical silicate
silDict = dustModel["Silicate"]
t = Template(**silDict)
colorList = ["r", "b", "m", "g", "y", "c"]
for nsil in silTList:
    tmplDir = modelDir + "{0}.dat".format(silNameList[nsil])
    modelSil = np.loadtxt(tmplDir)
    size_all  = modelSil[:, 0]
    wave_all  = modelSil[:, 1]
    kappa_all = modelSil[:, 2]
    sizeList = np.unique(size_all)
    for nsz in range(len(sizeList)): # [0, 79, 99]:
        sz = sizeList[nsz]
        fltr = size_all == sz
        wave  = wave_all[fltr]
        kappa = kappa_all[fltr]
        par = [nsil, sz-0.01]
        kappa_t = t(wave, par)
        color = colorList[nsz%len(colorList)]
        plt.plot(wave, kappa, linewidth=0.1, linestyle="-", color=color, label="{0}".format(sz))
        plt.plot(wave, kappa_t, linewidth=0.1, linestyle=":", color=color)
plt.xscale("log")
plt.yscale("log")
#plt.legend(loc="best", fontsize=10)
plt.title("{0}".format("Silicate"), fontsize=24)
#plt.savefig("dust_sil.png")
plt.show()
#"""

#"""
#For the amorphous oliven
graDict = dustModel["Graphite"]
t = Template(**graDict)
colorList = ["r", "b", "m", "g", "y", "c"]
for ngra in graTList:
    tmplDir = modelDir + "{0}.dat".format(graNameList[ngra])
    modelGra = np.loadtxt(tmplDir)
    size_all  = modelGra[:, 0]
    wave_all  = modelGra[:, 1]
    kappa_all = modelGra[:, 2]
    sizeList = np.unique(size_all)
    for nsz in range(len(sizeList)): #[0, 79, 99]:
        sz = sizeList[nsz]
        fltr = size_all == sz
        wave  = wave_all[fltr]
        kappa = kappa_all[fltr]
        par = [nsil, sz-0.01]
        kappa_t = t(wave, par)
        color = colorList[nsz%len(colorList)]
        plt.plot(wave, kappa, linewidth=0.1, linestyle="-", color=color, label="{0}".format(sz))
        plt.plot(wave, kappa_t, linewidth=0.1, linestyle=":", color=color)
plt.xscale("log")
plt.yscale("log")
#plt.legend(loc="best", fontsize=10)
plt.title("{0}".format("Graphite"), fontsize=24)
#plt.savefig("dust_sil.png")
plt.show()
#"""
