# coding: utf-8

""" From the mockstream runs, find the best mockstream that fits the data """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
import sys

# Third-party
from astropy import log as logger
import astropy.coordinates as coord
import astropy.units as u
import matplotlib.pyplot as pl
import numpy as np

# Custom
import gala.coordinates as gc
import gala.dynamics as gd
import gala.integrate as gi
import gala.potential as gp
from gala.observation import distance_modulus
from gala.units import galactic
from scipy.misc import logsumexp

import ophiuchus.potential as op
from ophiuchus.data import OphiuchusData
from ophiuchus.util import integrate_forward_backward
from ophiuchus.coordinates import Ophiuchus
from ophiuchus import galactocentric_frame, vcirc, vlsr, RESULTSPATH
from ophiuchus.experiments import MockStreamGrid
from ophiuchus.plot import plot_data_stream, surface_density

def ln_likelihood(model, ophdata):
    """
    Compute the likelihood by approximating the model density using a
    KDE of the points.

    Parameters
    ----------
    model : :class:`gala.dynamics.CartesianPhaseSpacePosition`
    ophdata : `ophiuchus.data.OphiuchusData`
    """

    # sum over model points, product over data points
    gal,vel = model.to_frame(coord.Galactic, galactocentric_frame=galactocentric_frame,
                             vcirc=vcirc, vlsr=vlsr)

    # bandwidth parameters for the KDE
    h_l = 0.1*u.deg
    h_d = 0.2*u.kpc
    h_vr = 0.5*u.km/u.s

    # get model coordinates
    model_l, model_b, model_d = gal.l.radian, gal.b.radian, gal.distance.to(u.kpc).value
    model_vr = vel[2].decompose(galactic).value

    # get data coordinates
    data_l, data_b, data_d = ophdata.coord.l.radian,ophdata.coord.b.radian,ophdata.coord.distance.to(u.kpc).value
    data_vr = ophdata.veloc['vr'].decompose(galactic).value

    # variances
    var_l = var_b = np.atleast_1d(h_l.decompose(galactic).value)**2
    var_d = (ophdata.coord_err['distance']**2 + h_d**2).decompose(galactic).value # extra distance spread