%matplotlib inline
import localgroup
import triangle
import sklearn
from sklearn import mixture
import numpy as np
import pickle
import matplotlib.patches as mpatches

"""
Explanation: Local Group Halo Properties: Demo Inference
We approximate the local group distance, radial velocity and proper motion likelihood function by sampling from the posterior distributions for these variables reported in the literature and transforming to kinematic variables in the M31-centric coordinate system.
End of explanation
"""


L = localgroup.Likelihood(isPair=True)
L.generate(Nsamples=200000)

L.set_PDF(mixture.GMM(n_components=10, covariance_type='full'))
L.approximate()

figure_obs = L.plot_samples(10, color='b', overlay=False)

"""
Explanation: Inside the Likelihood object is a "triplet" object called T, which contains an array of sample local groups, each with kinematic parameters consistent with the observational data. Let's plot these kinematic parameters in a "triangle" figure, to show all their 1 and 2-D marginal distributions.
End of explanation
"""


figure_model = L.model_gof(L.T.Nsamples, color="r", fig=None)

L.model_gof(L.T.Nsamples, color="r", fig=figure_obs)
red_patch = mpatches.Patch(color='red')
blue_patch = mpatches.Patch(color='blue')
figure_obs.legend(handles=[red_patch, blue_patch], labels=["Model Generated", "Observation Generated"])

figure_obs

"""
Explanation: The above plot shows a Gaussian Mixture model fitted Gaussians. The shaded regions show two standard deviations. The samples data has been preprocessed to zero the mean and scale by standard deviation. Since we are using the Gaussian Mixture Model to model the underlying PDF of the data, more components is always better. 
How to evaluate goodness of fit:
Due to lack of a standard goodness of fit test for GMM's, the best we can do is graphically show that the model reproduces the data well. We proceed by drawing a set of points from the fitted model, where each point is a local group with (MW_D, MW_vr, MW_vt, M33_D, M33_vr, M33_vt). We then plot the 1D and 2D marginalizations of the drawn point set and show that the marginalizations match the marginalizations of the true data.
End of explanation
"""


#name = "likelihood_fit.png"
#figure_obs.savefig("/afs/slac.stanford.edu/u/ki/mwillia1/Thesis/LocalGroupHaloProps/doc/thesis/plots/final_plots/Pair_LGMM_GOF.pdf", dpi=1200)

"""
Explanation: The above plot shows that the points drawn from the model create a population that is very similar to the true data.
End of explanation
"""


path = '/lustre/ki/pfs/mwillia1/LG_project/Consuelo_Boxes/4001/quad_dat_M31_larger_bgc2.npy'
#path = '/afs/slac.stanford.edu/u/ki/mwillia1/Thesis/data_files/MW_M31_pairs.txt'
npoints = 5000
halo_props = ['MW_Mvir', 'M31_Mvir', 'M33_Mvir']

Tr = localgroup.Triplet(isPair=True)
Tr.read_sim_points(path, npoints, halo_props, h=0.7, a=1.0, npy=True)

Tr.transform_to_M31(sim=True)

#Tr.mass_filter('sim')

Tr.dist_filter((Tr.sim_samples[:,0] < 1))

Tr.preprocess(L.samples_means, L.samples_stds, mode='sim')

sim_plot = Tr.plot_kinematics('sim', L.samples_means, L.samples_stds, color='c', fig=None)

Tr.plot_kinematics('sim', L.samples_means, L.samples_stds, color='g', fig=figure_model)

red_patch = mpatches.Patch(color='red')
green_patch = mpatches.Patch(color='green')
figure_model.legend(handles=[red_patch, green_patch], labels=["Model Generated", "Consuelo Triplets"])
figure_model

#name = 'trips_000_054_like_prior.png'
#sim_plot.savefig('/afs/slac.stanford.edu/u/ki/mwillia1/Thesis/LocalGroupHaloProps/doc/thesis/plots/asurps/'+name)

Tr.compute_model_weights(L, "sim")

N95 = Tr.calculate_N95()
N95

MW_C200 = Tr.calculate_cvir_to_c200(Tr.MW.Cvir)
M31_C200 = Tr.calculate_cvir_to_c200(Tr.M31.Cvir)
#M33_C200 = Tr.calculate_cvir_to_c200(Tr.M33.Cvir)

MW_M200 = Tr.calculate_Mvir_to_M200(Tr.MW.Mvir, Tr.MW.Cvir, MW_C200)
M31_M200 = Tr.calculate_Mvir_to_M200(Tr.M31.Mvir, Tr.M31.Cvir, M31_C200)
#M33_M200 = Tr.calculate_Mvir_to_M200(Tr.M33.Mvir, Tr.M33.Cvir, M33_C200)

total_200 = M31_M200+MW_M200

total = Tr.MW.Mvir + Tr.M31.Mvir
all_mvir = np.log10(np.transpose(np.vstack((Tr.MW.Mvir, Tr.M31.Mvir, total))))
labs = ["$M_{MW}$", "$M_{M31}$", "$M_{LG}$"]
#total = MW_M200 + M31_M200
#all_mvir = np.log10(np.transpose(np.vstack((MW_M200, M31_M200, total))))
#labs = ["MW Mvir", "M31 Mvir", "MW+M31"]

all_mvir = np.transpose(np.vstack((Tr.MW.Cvir, Tr.M31.Cvir)))
labs = ["$Cvir_{MW}$", "$Cvir_{M31}$"]

all_mvir.shape

figure = triangle.corner(all_mvir, labels=labs, quantiles=[0.16,0.5,0.84], fig=None, weights=np.transpose(Tr.weights),\
                         plot_contours=True, show_titles=True, title_args={"fontsize": 16}, label_args={"fontsize": 16},\
                         plot_datapoints=False, bins=20, color='k')

h = np.load('/lustre/ki/pfs/mwillia1/LG_project/Consuelo_Boxes/4002/4002hlist.npy')

mw_h = h[np.abs(np.log10(h['mvir'])-11.7) < .1]
m31_h = h[np.abs(np.log10(h['mvir'])-12) < .1]
m33_h = h[np.abs(np.log10(h['mvir'])-10.8) < .1]

mwc = mw_h['rvir']/mw_h['rs']
m31c = m31_h['rvir']/m31_h['rs']
m33c = m33_h['rvir']/m33_h['rs']

mwc.shape, m31c.shape, m33c.shape

mwc = mwc[0:100000]
m31c = m31c[0:100000]
m33c = m33c[0:100000]

all_mvir = np.transpose(np.vstack((mwc, m31c)))
labs = ["$Cvir_{MW}$", "$Cvir_{M31}$", "$Cvir_{M33}$"]

figure_c = triangle.corner(all_mvir, labels=labs, quantiles=[0.,0.5,0.84], fig=figure, weights=None,\
                         plot_contours=True, show_titles=True, title_args={"fontsize": 16}, label_args={"fontsize": 16},\
                         plot_datapoints=False, bins=20, color='k')

figure2 = triangle.corner(all_mvir, labels=labs, quantiles=[0.16,0.5,0.84], fig=figure, weights=Tr.weights,\
                         plot_contours=True, show_titles=False, title_args={"fontsize": 12}, \
                         plot_datapoints=False, bins=20, color='r')

figure

red_patch = mpatches.Patch(color='red')
green_patch = mpatches.Patch(color='k')
figure.legend(handles=[red_patch, green_patch], labels=["Box 4001 bgc2", "Box 4001 hlist"])
figure

"""
Explanation: Reading Simulation Points:
Below we read the preconfigured files containing the Consuelo (soon to be Dark Sky) Local Group analogs into a Triplet object. We plot the marginalizations of the simulation data, which allows us to compare with the LG prior.
End of explanation
"""


Tr.weights.shape

xdata = total
plot(xdata, Tr.weights)
xscale('log')

bins = 20
xbins = np.logspace(np.log10(xdata.min()), np.log10(xdata.max()), bins)
ybins = [np.sum(Tr.weights[(xdata[:]>xbins[i]) & (xdata[:]<xbins[i+1])]) for i in range(bins-1)]
print len(xbins)
print len(ybins)
print len(xbins[:bins-1])

plot(xbins[:bins-1], ybins)
xscale('log')

xtrip = np.copy(xbins[:bins-1])
ytrip = np.copy(ybins)

xtripmw = np.copy(xbins[:bins-1])
ytripmw = np.copy(ybins)

xpairmw = np.copy(xbins[:bins-1])
ypairmw = np.copy(ybins)

xpair = np.copy(xbins[:bins-1])
ypair = np.copy(ybins)

plot(xtripmw, ytripmw, color='r')
plot(xpairmw, ypairmw, color='b')
ylabel('P(MMW)')
xlabel('MMW')

red_patch = mpatches.Patch(color='red')
blue_patch = mpatches.Patch(color='blue')
legend(["Trips, M33 dynamics", "Pairs, M33 existence"])
xscale('log')
#savefig("/afs/slac.stanford.edu/u/ki/mwillia1/Thesis/LocalGroupHaloProps/doc/thesis/plots/pdf_mw_plot.png")

"""
Explanation: 1D posterior work
End of explanation
"""


