from __future__ import print_function, division

import sys
import os, os.path
import pandas as pd
import numpy as np

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

period_rng = (50, 300)
rp_rng = (0.75, 20)

# Read synthetic catalog
koi_file = sys.argv[1]

kois = pd.read_hdf(koi_file, 'kois')
theta_true = list(pd.read_hdf(koi_file, 'theta'))


# Load completeness contour
completeness_file = 'completeness.npz'

d = np.load(completeness_file)

comp = d['comp']
period_grid = d['period_grid']
rp_grid = d['rp_grid']
comp_inds = d['inds']  #indices of stellar table used

# A double power law model for the population.
def population_model(theta, period, rp):
    lnf0, beta, alpha = theta
    v = np.exp(lnf0) * np.ones_like(period)
    for x, rng, n in zip((period, rp),
                         (period_rng, rp_rng),
                         (beta, alpha)):
        n1 = n + 1
        v *= x**n*n1 / (rng[1]**n1-rng[0]**n1)
    return v

# The ln-likelihood function given at the top of this post.
koi_periods = np.array(kois.koi_period)
koi_rps = np.array(kois.koi_prad)
vol = np.diff(period_grid, axis=0)[:, :-1] * np.diff(rp_grid, axis=1)[:-1, :]
def lnlike(theta):
    pop = population_model(theta, period_grid, rp_grid) * comp
    pop = 0.5 * (pop[:-1, :-1] + pop[1:, 1:])
    norm = np.sum(pop * vol)
    ll = np.sum(np.log(population_model(theta, koi_periods, koi_rps))) - norm
    return ll if np.isfinite(ll) else -np.inf

# The ln-probability function is just propotional to the ln-likelihood
# since we're assuming uniform priors.
bounds = [(-5, 5), (-5, 5), (-5, 5)]
def lnprob(theta):
    # Broad uniform priors.
    for t, rng in zip(theta, bounds):
        if not rng[0] < t < rng[1]:
            return -np.inf
    return lnlike(theta)

# The negative ln-likelihood is useful for optimization.
# Optimizers want to *minimize* your function.
def nll(theta):
    ll = lnlike(theta)
    return -ll if np.isfinite(ll) else 1e15



# Optimize, and then run the chain
from scipy.optimize import minimize
theta_0 = np.array([-0.3, -1.5, -0.8])
r = minimize(nll, theta_0, method="L-BFGS-B", bounds=bounds)

import emcee

ndim, nwalkers = len(r.x), 16
pos = [r.x + 1e-5 * np.random.randn(ndim) for i in range(nwalkers)]
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob)

# Burn in.
pos, _, _ = sampler.run_mcmc(pos, 1000)
sampler.reset()

# Production.
pos, _, _ = sampler.run_mcmc(pos, 4000)


import corner
corner.corner(sampler.flatchain, labels=[r"$\ln F$", r"$\beta$", r"$\alpha$"],
                truths=(theta_true[0], theta_true[2], theta_true[1]))

filebase = os.path.splitext(koi_file)[0]
plt.savefig('{}_corner.png'.format(filebase))
                
np.save('{}_chains.npy'.format(filebase), sampler.flatchain)

