#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Plots the genealogical tree generated by a bootstrap filter. Model is 

X_t|X_{t-1}=x_{t-1} ~ N(mu+phi(x_{t-1}-mu), sigma^2)
Y_t|X_t=x_t ~ Poisson(exp(x_t))

See Figure 11.2 for more details. (Typo: the book says rho=0.5 
while we took rho=0.9)

"""

from __future__ import division, print_function

from matplotlib import pyplot as plt
import numpy as np
# import seaborn as sb

import particles
from particles import distributions as dists
from particles import state_space_models

# set up models, simulate and save data
T = 100
mu0 = 0.
phi0 = 0.9
sigma0 = .5  # true parameters
ssm = state_space_models.DiscreteCox(mu=mu0, phi=phi0, sigma=sigma0)
true_states, data = ssm.simulate(T)
fkmod = state_space_models.Bootstrap(ssm=ssm, data=data)

# run particle filter, compute trajectories 
N = 100
pf = particles.SMC(fk=fkmod, N=N, store_history=True)
pf.run()
Bs = pf.hist.compute_trajectories()

# PLOT
# ====
# sb.set_palette("dark")
plt.style.use('ggplot')
savefigs = True  # False if you don't want to save plots as pdfs

plt.figure()
plt.xlabel('t')
for n in range(N):
    plt.plot(Bs[:, n], 'k')
if savefigs:
    plt.savefig('genealogy.pdf')

plt.show()
