from __future__ import (absolute_import, division, print_function)
from lamost import LamostDataset
from cannon.model import CannonModel
from cannon.spectral_model import draw_spectra, diagnostics, triangle_pixels, overlay_spectra, residuals
import numpy as np
import pickle
from matplotlib import rc

plt.rc('text', usetex=True)
plt.rc('font', family='serif')

tr_files = np.genfromtxt("example_LAMOST/Training_Data.txt", dtype=str)
test_files = np.loadtxt("example_LAMOST/Test_Data.txt", dtype=str)
dataset = LamostDataset("example_LAMOST/Data_All",
                        tr_files, test_files,
                        "example_DR12/reference_labels.csv")

# Make a nice plot of a typical high-S/N spectrum
# with a continuum fit

# SNR of jj=0 is 134.13631569901949
# SNR of jj=150 is 172.30452688124373
# SNR of jj=10 is 280.99611610353145

tr_cont, test_cont = pickle.load(open("cont.p", "r"))

jj = 0 
bad1 = dataset.tr_fluxes[jj,:] == 0.
bad2 = np.var(dataset.tr_ivars, axis=0) == 0.
bad = np.logical_or(bad1, bad2)
x = dataset.wl