"""
File: 02-dynamic_solution.py

Michel Bierlaire
Thu Aug 07 2025, 17:43:46
"""

import biogeme.biogeme_logging as blog
import pandas as pd
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.database import Database
from biogeme.expressions import (
    Beta,
    Draws,
    Expression,
    MonteCarlo,
    PanelLikelihoodTrajectory,
    Variable,
    log,
)
from biogeme.models import logit, loglogit
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)

# As the estimation time may be long, we ask Biogeme to report the details of the iterations.
logger = blog.get_screen_logger(level=blog.INFO)

# **Tip:**<div class="alert alert-block alert-info">It is advised to start working with a low number of draws, until
# the script is working well. Then, increase the number of draws to 10000, say. Then, execute the script overnight.
# </div>
number_of_draws = 10


# # Dynamic Choice Models

# We analyze the smoking behavior of individuals, as a function of their age and the price of tobacco using synthetic
# data. We develop a model that predicts, for every year, the probability to smoke or not.

# ## Data

# We use synthetic data that has been generated as follows. We postulate a true model for the data generation process.
# It is a mixture of logit models with
# two alternatives: ``smoking`` or ``not smoking``.
# The utility for individual $n$ associated with "not smoking" in year $t$ is
# \begin{equation}
# U_{0nt}= \varepsilon_{0nt}
# \end{equation}
#  and the utility associated with "smoking" is
# \begin{equation}
# U_{1nt}= \beta_{nt} y_{n,t-1} + \beta^p_{nt} P_{t} + c_n + \varepsilon_{1nt},
# \end{equation}
# where
#
# - $\beta_{nt} = 10$,
#
# - $y_{n,t-1}=1$ if $n$ is smoking at time $t-1$, $0$ otherwise,
#
# - $\beta^p_{nt} = -0.1$,
#
# - $P_t$ is the price of cigarettes at time $t$,
#
# - $c_n$ is an individual specific constant that captures the a priori, intrinsic attraction of each individual
# towards smoking. It is assumed to be normally distributed in the population, with zero mean and standard deviation
# 50: $N(0, 50^2)$, and constant over $t$.
#
# We generate a sample of 1000 individuals, and we simulate their smoking behavior between the age of 16 until the age
# of 100.
#
# The date of birth of each individual is uniformly distributed between 2000 and 2020.
# The price of cigarettes in 2000 is supposed to be 10. The price of cigarettes in year $t$ is
# $$P_t = 10 \cdot 1.02^{t-2000},$$
# which represents a price increase of 2% per year.


# ## True value of the parameters

# We store the true value of the parameters for future comparison
true_parameters = pd.DataFrame(
    {'Value': [-0.1, 10, 0, 50]},
    index=['coef_price', 'beta_last_year', 'cte_mean', 'cte_std'],
)


# ## Data

# The observations are available in the following data file.
df = pd.read_table('smoking.dat', sep=',')
display(df)


# The data contains the following columns:
#
# - the age of the individual,
# - the price of the cigarettes,
# - a variable that is 1 if the individual is smoking, 0 otherwise,
# - a variable that is 1 if the individual was smoking last year, 0 otherwise,
# - a unique id for each individual,
# - a variable that is 1 if the individual was smoking at the age of 45, in the beginning of the observation period.
database = Database('smoking', df)


# Variables
Price = Variable('Price')
Smoking = Variable('Smoking')
LastYear = Variable('LastYear')


# ## Estimation procedure


# The following procedure estimates the choice model (or read the estimation results from file), and returns the
# estimated parameters in a Pandas format.
def estimate(
    the_logprob: Expression, the_name: str, the_database: Database
) -> pd.DataFrame:
    """Estimates the choice model (or read the estimation results from
        file if recycle is True), and returns the estimated parameters
        in a Pandas format.

    :param the_logprob: formula for the log likelihood function
    :param the_name: name of the model (important for the output files)
    :param the_database: database
    :return: estimated values of the parameters and statistics.
    """
    biogeme = BIOGEME(
        the_database,
        the_logprob,
        number_of_draws=number_of_draws,
    )
    biogeme.model_name = the_name
    print(f'Sample size: {biogeme.sample_size}')
    print(f'Number of observations: {biogeme.number_of_observations}')

    results: EstimationResults = biogeme.estimate()
    print(results.short_summary())
    pandas_results = get_pandas_estimated_parameters(estimation_results=results)
    return pandas_results.set_index('Name')


# ## Static model

# The static model considers the data as cross-sectional. No state dependence, and no serial correlation is captured.
cte_mean = Beta('cte_mean', 0, None, None, 0)
coef_price = Beta('coef_price', 0, None, None, 0)


# Model specification
v_smoking = coef_price * Price + cte_mean
v_non_smoking = 0
v = {0: v_non_smoking, 1: v_smoking}
log_probability = loglogit(v, None, Smoking)


# Estimation
r_static = estimate(log_probability, 'static_model', database)
display(r_static)

# ## Comparison of the estimates
summary = pd.concat([true_parameters['Value'], r_static['Value']], axis='columns')
summary.columns = ['True', 'Static']
summary.fillna('')
display(summary)

# The estimated price coefficient is not significant. Indeed, price is the only variable that the model considers.
# Ignoring state dependence generates endogeneity. The model "thinks" that individuals are insensitive to price,
# as they choose an alternative that is expensive.

# ## Dynamic model

# The dynamic model adds the choice of last year as an explanatory variable
beta_last_year = Beta('beta_last_year', 0, None, None, 0)


# Model specification
v_smoking = beta_last_year * LastYear + coef_price * Price + cte_mean
v_non_smoking = 0
v = {0: v_non_smoking, 1: v_smoking}
log_probability = loglogit(v, None, Smoking)


# Estimation
r_dynamic = estimate(log_probability, 'dynamic_model', database)
display(r_dynamic)


# ### Comparison of the estimates
summary = pd.concat(
    [
        true_parameters['Value'],
        r_static['Value'],
        r_dynamic['Value'],
    ],
    axis='columns',
)
summary.columns = ['True', 'Static', 'Dynamic']
summary.fillna('')
display(summary)


# The introduction of the lag variable has increased a lot the final log likelihood from `-57529.39` to `-2226.1`.
# Note that the error term in the model is not the same as in the true model. Indeed, serial correlation has been
# ignored. Therefore, the coefficients cannot be directly compared. But their ratio can be compared, as it cancels the
# scale. And it can be seen that the ratios are almost the same.
ratio_true = summary.loc['coef_price', 'True'] / summary.loc['beta_last_year', 'True']
print(f'Ratio coef_price / beta_last_year (True): {ratio_true:.3g}')
ratio_dynamic = (
    summary.loc['coef_price', 'Dynamic'] / summary.loc['beta_last_year', 'Dynamic']
)
print(f'Ratio coef_price / beta_last_year (Dynamic): {ratio_dynamic:.3g}')


# ## Static model with serial correlation

# We now introduce the agent effect to capture serial correlation. First, we tell Biogeme that the data is organized
# as a panel, meaning that there are several observations corresponding to the same individuals.
# Therefore, instead of considering that there is a sample of 11000 independent observations, Biogeme knows that there
# is actually a sample of 1000 individuals, for which a trajectory is observed.


# Declaring the panel nature of the data
database.panel('Id')
print(
    f'Sample size, accounting for the panel nature of the data: {database.num_rows()}'
)


# Model specification
cte_std = Beta('cte_std', 10, None, None, 0)
cte = cte_mean + cte_std * Draws('agent', 'NORMAL_ANTI')
v_smoking = coef_price * Price + cte
v_non_smoking = 0
v = {0: v_non_smoking, 1: v_smoking}
conditional_observation_probability = logit(v, None, Smoking)
conditional_individual_probability = PanelLikelihoodTrajectory(
    conditional_observation_probability
)
log_probability = log(MonteCarlo(conditional_individual_probability))


# Estimation
r_serial_static = estimate(log_probability, 'static_model_serial', database)
display(r_serial_static)


# ## Dynamic model with serial correlation

# We now introduce the state dependence in the model, to make it dynamic.

# Model specification
v_smoking = beta_last_year * LastYear + coef_price * Price + cte
v_non_smoking = 0
v = {0: v_non_smoking, 1: v_smoking}
conditional_observation_probability = logit(v, None, Smoking)
conditional_individual_probability = PanelLikelihoodTrajectory(
    conditional_observation_probability
)
log_probability = log(MonteCarlo(conditional_individual_probability))


# Estimation
r_serial_dynamic = estimate(log_probability, 'dynamic_model_serial', database)
display(r_serial_dynamic)


# ### Comparison of the estimates
summary = pd.concat(
    [
        true_parameters['Value'],
        r_static['Value'],
        r_dynamic['Value'],
        r_serial_static['Value'],
        r_serial_dynamic['Value'],
    ],
    axis='columns',
)
summary.columns = ['True', 'Static', 'Dynamic', 'Static + serial', 'Dynamic + serial']
summary.fillna('')
display(summary)


# If the number of draws is sufficiently high, we observe that the parameters are quite well recovered. We can
# actually test the hypothesis that they are equal to their true value, using a $t$-test.
def t_test(param: str, true_value: float) -> float:
    """Calculates the t-statistic to test the hypothesis that the true
        value of the parameter is indeed the true one.

    :param param: name of the parameter
    :param true_value: true value of the parameter
    :return: t-statistic
    """
    return (r_serial_dynamic.loc[param, 'Value'] - true_value) / r_serial_dynamic.loc[
        param, 'Robust std err.'
    ]


# `coef_price`
the_test = t_test('coef_price', -0.1)
print(f'Test for coef_price: {the_test:.3g}')

# `beta_last_year`
the_test = t_test('beta_last_year', 10)
print(f'Test for beta_last_year: {the_test:.3g}')


# `cte_mean`
the_test = t_test('cte_mean', 0)
print(f'Test for cte_mean: {the_test:.3g}')


# `cte_std`
the_test = t_test('cte_std', 50)
print(f'Test for cte_std: {the_test:.3g}')

# Each of this $t$-test is sufficiently low so that the hypothesis cannot be rejected.