%matplotlib inline
import numpy as np
from sklearn.cluster import KMeans
from astroML.stats import binned_statistic
import matplotlib.pyplot as plt

"""
Explanation: This notebook explores the PHAT v2 artificial star test (AST) results, and how to use them in m31hst.
End of explanation
"""


from m31hst.phatast import load_phat_ast_table
t = load_phat_ast_table()

"""
Explanation: Assuming that the Williams et al 2014 Table 6 file was downloaded to the correct location (see the README of m31hst), the AST result catalog can be loaded with the load_phat_ast_table into an astropy table.
End of explanation
"""


km = KMeans(n_clusters=6)
xy = np.vstack((t['ra'], t['dec'])).T
km.fit(xy)
centers = km.cluster_centers_
print centers

colors = ['c', 'm', 'y', 'k', 'r', 'b']

figure = plt.figure(figsize=(6, 6))
ax = figure.add_subplot(111)
ax.scatter(t['ra'], t['dec'],
           marker='.', edgecolors='None', s=1,
           c=km.labels_, rasterized=True)
ax.scatter(centers[:, 0], centers[:, 1],
           marker='x', s=20, c='None',
           edgecolors=colors, zorder=10)
ax.set_xlim(11.6, 10.6)
figure.show()

"""
Explanation: PHAT did AST work in six fields, but unfortunately they don't label what field a star belongs to. This labelling is important because it connects the AST results to measurements of local stellar density. Here I'm showing that we can use K-Means clustering on the stars to reliably label stars into the six AST fields.
End of explanation
"""


f2 = plt.figure(figsize=(6, 6))
ax = f2.add_subplot(111)
for label, c in zip(range(6), colors):
    sel = np.where(km.labels_ == label)[0]
    tt = t[sel]
    diffs = np.abs(tt['f475w_in'] - tt['f475w_out'])
    s = np.where(diffs < 10.)[0]
    err, edges = binned_statistic(tt['f475w_in'][s],
                                  diffs[s],
                                  statistic='median', bins=50)
    ax.scatter(tt['f475w_in'][::2], diffs[::2],
               s=1, marker='.', alpha=0.4,
               edgecolors='None', facecolors=c)
    ax.plot(edges[:-1], err, ls='-', c=c)
ax.set_xlabel(r'$m$')
ax.set_ylabel(r'$\Delta m$')
ax.set_ylim(0, 0.2)
ax.set_xlim(15., 25.)
f2.show()
        

"""
Explanation: For each AST field, let's plot the median photometric error in the F475W band as a function of luminosity:
End of explanation
"""


f3 = plt.figure(figsize=(6, 6))
ax = f3.add_subplot(111)
n_bins = 80
for label, c in zip(range(6), colors):
    sel = np.where(km.labels_ == label)[0]
    tt = t[sel]
    dropped = np.where(tt['f475w_out'] >= 90.)[0]
    drop_flag = np.ones(len(tt), dtype=np.float)
    drop_flag[dropped] = 0.
    total_recovered, edges = binned_statistic(tt['f475w_in'],
                                              drop_flag,
                                              statistic='sum',
                                              bins=n_bins)
                                              # range=np.array([15., 28.]))
    count, edges = binned_statistic(tt['f475w_in'],
                                    drop_flag,
                                    statistic='count',
                                    bins=n_bins)
                                    # range=np.array([15., 28.]))
    ax.plot(edges[:-1], total_recovered / count, ls='-', c=c)
ax.set_xlabel(r'$m$')
ax.set_ylabel(r'$\Delta m$')
ax.set_ylim(-0.05, 1.05)
ax.set_xlim(15., 28.)
f3.show()

"""
Explanation: We can also plot the completeness (fraction of artificial stars recovered successfully) in the F475W band as a function of luminosity for each AST field:
End of explanation
"""


from m31hst.phatast import PhatAstTable
tbl = PhatAstTable()
tbl.fields

"""
Explanation: We provide a PhatAstTable object to provide access to the AST data, sorted by field:
End of explanation
"""


tbl.write_crowdfile_for_field("crowdfile.txt", 0)

%%bash
head crowdfile.txt

"""
Explanation: Here we create a crowding file for use in StarFISH from stars in the outer-most (0) AST field:
End of explanation
"""


