from kepcatsim import simulate_survey, BURKE_QUERY

"""
Explanation: Let's generate a catalog of planets, including binary stars.  There are parameters controlling both the planet population and the binary star population:
$$ \theta = (F, \beta, \alpha, f_B, \gamma) $$
$F$, $\beta$ and $\alpha$ are the power-law parameters of the true planet population, $f_B$ is the binary fraction, and $\gamma$ is the mass-ratio power law index.  We want to generate a mock catalog according these parameters, with the Kepler stars as the starting point.
End of explanation
"""


BURKE_QUERY

theta_true = [-0.3, -0.8, -1.5, 0.5, 0.3] #lnF0, beta, alpha, fB, gamma
period_rng = (50,300)
rp_rng = (0.75,20)
mes_threshold = 10.
survey = simulate_survey(theta_true, query=BURKE_QUERY, 
                         period_rng=period_rng, rp_rng=rp_rng, mes_threshold=mes_threshold)

print(len(survey.planets))
survey.planets.head(10)

survey.planets.groupby('which').count()

"""
Explanation: simulate_survey has all the simulation magic in it. Here's what it does: 

Selects from the Kepler stellar catalog according to query, 
Simulates binary stars for each of those stars, according to the binary star parameters, 
Simulates planets according to planet parameters, and
Dilutes & re-interprets those signals according to the "single star only" hypothesis, keeping only those that are detected.

Here's the Burke query, just for a reminder:
End of explanation
"""


from kepcatsim.inference import DoublePowerLaw, ProbabilisticModel, run_mcmc
import corner
%matplotlib inline

popmodel = DoublePowerLaw(period_rng, rp_rng)
model = ProbabilisticModel(survey, popmodel)
sampler = run_mcmc(theta_true[:3], model)

"""
Explanation: That's a bit interesting--- this qualitatively matches what Fressin+ (2013) predict for the fraction of "false positives" from companion star + transiting planet.
Anyway, let's do some inference:
End of explanation
"""


corner.corner(sampler.flatchain, labels=[r"$\ln F$", r"$\beta$", r"$\alpha$"],
                truths=(theta_true[:3]));


"""
Explanation: See what's up:
End of explanation
"""


theta_true2 = [-0.3, -0.8, -1.5, 0., 0.3] #lnF0, beta, alpha, fB, gamma
survey2 = simulate_survey(theta_true2, query=BURKE_QUERY, 
                         period_rng=period_rng, rp_rng=rp_rng, mes_threshold=mes_threshold)
model2 = ProbabilisticModel(survey2, popmodel)
sampler2 = run_mcmc(theta_true2[:3], model2)

corner.corner(sampler2.flatchain, labels=[r"$\ln F$", r"$\beta$", r"$\alpha$"],
                truths=(theta_true2[:3]));

"""
Explanation: OK, great-- as a control, let's do this again, but setting $f_B$ to zero.
End of explanation
"""


from kepcatsim.inference import DoublePowerLaw, ProbabilisticModel, run_mcmc
import corner 

popmodel = DoublePowerLaw(period_rng, rp_rng)
theta_true3 = [-0.3, -0.8, -1.5, 0.5, 0.3] #lnF0, beta, alpha, fB, gamma
survey3 = simulate_survey(theta_true3, query=BURKE_QUERY, 
                         period_rng=period_rng, rp_rng=rp_rng, 
                          mes_threshold=mes_threshold, per_star=True)
model3 = ProbabilisticModel(survey3, popmodel)
sampler3 = run_mcmc(theta_true3[:3], model3)

survey3.planets.groupby('which').count()

corner.corner(sampler3.flatchain, labels=[r"$\ln F$", r"$\beta$", r"$\alpha$"],
                truths=(theta_true3[:3]));

"""
Explanation: OK, now let's see what happens if we interpret $F$ as the per star occurrence rate instead of per system:
End of explanation
"""


