%matplotlib inline
from astropy.table import Table
import numpy as np
import matplotlib.pyplot as plt
data = Table.read('GTR-ADM-QSO-ir-testhighz_findbw_lup_2016_starclean.fits')

# X is in the format need for all of the sklearn tools, it just has the colors
X = np.vstack([ data['ug'], data['gr'], data['ri'], data['iz'], data['zs1'], data['s1s2']]).T
y = np.array(data['labels'])

from sklearn.ensemble import RandomForestClassifier
clf = RandomForestClassifier(n_estimators=10, max_depth=15, min_samples_split=2, n_jobs=-1, random_state=42)
clf.fit(X, y)

"""
Explanation: Final SpIES High-z Quasar Selection
Notebook performing selection of $3.5<z<5$ quasars from SDSS+SpIES data.
First load the training data, then instantiate and train the algorithm; see https://github.com/gtrichards/QuasarSelection/blob/master/SpIESHighzCandidateSelection.ipynb
End of explanation
"""


data2 = Table.read('GTR-ADM-QSO-ir_good_test_2016.fits')
#data2 = Table.read('GTR-ADM-QSO-ir_classifiers_good_test_2016.fits')
#data2 = Table.read('gtr-adm-qso-ir_photoz_in7_good_test_2016.fits')

Xtest = np.vstack([ data2['ug'], data2['gr'], data2['ri'], data2['iz'], data2['zs1'], data2['s1s2']]).T

print len(Xtest),Xtest[:5]

"""
Explanation: Second, load the test data
Test Data
Test set data set was made as follows (see 18 April 2016 README entry):
maketest_2016.py

Output is:
classifiers_out = open('GTR-ADM-QSO-ir_classifiers_good_test_2016.dat','w') 
others_out= open('GTR-ADM-QSO-ir_others_good_test_2016.dat','w')
czr_out = open('GTR-ADM-QSO-ir_photoz_in7_good_test_2016.dat','w')

Really need the first two files combined (so that we have both RA/Dec and colors in one place).
But couldn't merge them with TOPCAT or STILTS.  So had to break them into 3 pieces (with TOPCAT),
then used combine_test_files_STILTS.py to merge them together (just changing the input/output file names by hand).  
Actually ran this on dirac so that I'd have more memory than on quasar.  Copied the output files back to quasar and merged them together with TOPCAT.
So<br>
GTR-ADM-QSO-ir_others_good_test_2016a.dat + GTR-ADM-QSO-ir_classifiers_good_test_2016a.dat<br>
gives<br>
GTR-ADM-QSO-ir_good_test_2016a.dat<br>
(and so on for "b" and "c").
Then<br>
GTR-ADM-QSO-ir_good_test_2016a.dat + GTR-ADM-QSO-ir_good_test_2016b.dat + GTR-ADM-QSO-ir_good_test_2016c.dat<br>
gives<br>
GTR-ADM-QSO-ir_good_test_2016.dat
N.B.  Had to delete the intermediate files from quasar to save disk space.
GTR: Re-did this with combine_test_files_STILTSn.py (on quasar) since I wanted to include the i-band and ch2 magnitudes.  The output file is still GTR-ADM-QSO-ir_good_test_2016.dat.
Now read in the test file and convert it to an appropriate array format for sklearn.
End of explanation
"""


ypred = clf.predict(Xtest)

print ypred[:5]
print len(ypred)

data2['ypred'] = ypred
data2.write('GTR-ADM-QSO-ir_good_test_2016_out.fits', format='fits')

"""
Explanation: Quasar Candidates
Finally, do the classification and output the test file, including the predicted labels.
End of explanation
"""


mask = (data2['ypred']==0)
quasars = data2[mask]
print len(quasars)

"""
Explanation: This gives about 27,000 candidates in all and about 3000 on Stripe 82.
You can get the output file from
http://www.physics.drexel.edu/~gtr/outgoing/spies/GTR-ADM-QSO-ir_good_test_2016_out.fits.bz2
End of explanation
"""


from sklearn.preprocessing import StandardScaler 
scaler = StandardScaler()
scaler.fit(X)  # Use the full training set now
XS = scaler.transform(X)
XStest = scaler.transform(Xtest)

from sklearn.ensemble import BaggingClassifier
from sklearn.neighbors import KNeighborsClassifier
bag = BaggingClassifier(KNeighborsClassifier(n_neighbors=5), max_samples=0.5, max_features=1.0)

bag.fit(XS, y)

from dask import compute, delayed

def processBAG(Xin):
    return bag.predict(Xin)

# Create dask objects
# Reshape is necessary because the format of x as drawm from Xtest 
# is not what sklearn wants.
dobjsBAG = [delayed(processBAG)(x.reshape(1,-1)) for x in XStest]
#print dobjs

"""
Explanation: The main test set classification ends here.  But see below for some experimenting with bagging.
Also, see<br>
SpIESHighzQuasarsS82all.py<br>
for some stand-alone code that classifies Stripe 82 candidates by RF, SVM and bagging.

Bagging classifier
If we want to try bagging instead, here is how we'd do that.  But this takes a really long time unless we restrict ourselves to just Stripe 82.  We'll use dask to schedule the processes, otherwise it kills the kernel.
End of explanation
"""


print dobjsBAG[:5]
dobjsBAG1 = dobjsBAG[:1000000]

import dask.threaded
ypredBAG1 = compute(*dobjsBAG1, get=dask.threaded.get)

ypredBAG1 = np.array(ypredBAG1).reshape(1,-1)[0]

print ypredBAG1[:100]

"""
Explanation: See if we can even handle the first million objects, let alone the full data set.
End of explanation
"""


data2['ypred'] = ypredBAG
data2.write('GTR-ADM-QSO-ir_good_test_2016_out_bag.fits', format='fits')

"""
Explanation: IFF we could get this to work, then we'd produce output as follows.  But this always crashes.  Tried running on dirac and it crashes there too.  So, had to concentrate on Stripe 82.  See SpIESHighzQuasarsS82all.py.
End of explanation
"""


