"""
File: 04-wooldridge_solution.py

Michel Bierlaire
Sun Aug 10 2025, 17:31:54
"""

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
from biogeme.results_processing import 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 with Panel Effects

# We analyze again 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 postulate a true model for the data generation process. It is a mixture of logit models where the utility
# associated with "not smoking" 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 cigarets 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)$,

# ## True value of the parameters

true_parameters = pd.DataFrame(
    {'Value': [-0.1, 10, 0, 50]},
    index=['coef_price', 'beta_last_year', 'cte_mean', 'cte_std'],
)


# ## Data

# We observe every individual only from the age of 45 and the age of 55.
df = pd.read_table('smoking55.dat', sep=',')
display(df)


# Different values for age
print(df['Age'].unique())


# 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('smoking55', df)


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


# We declare that the data set contains panel data.
database.panel('Id')


# ## Estimation procedure


# The following procedure estimates the choice model, and returns the estimated parameters in a Pandas format. If the
# model happens to have been already estimated, the estimation results are read from the pickle file and reported.
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
    results = biogeme.estimate()
    print(results.short_summary())
    pandas_results = get_pandas_estimated_parameters(estimation_results=results)
    return pandas_results.set_index('Name')


# ## Dynamic model with serial correlation

# In the previous quiz, we have estimated a dynamic model with panel effects to account for serial correlation.
cte_mean = Beta('cte_mean', 0, None, None, 0)
cte_std = Beta('cte_std', 1, None, None, 0)
cte = cte_mean + cte_std * Draws('agent', 'NORMAL_ANTI')
coef_price = Beta('coef_price', 0, None, None, 0)
beta_last_year = Beta('beta_last_year', 0, None, None, 0)


# 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_serial_dynamic['Value'],
    ],
    axis=1,
)
summary.columns = ['True', 'Dynamic + serial']
summary.fillna('')
display(summary)

# We observe here the issue of the "initial condition problem". Although the model specification is correct (it is the
# same model as the data generation process), the values of the parameters are not correctly recovered. It is because
# the first observed choice, that is, the fact that an individual is smoking at the age of 45, is strongly correlated
# with the agent effect. This creates endogeneity. One visible consequence is the positive price coefficient. We are
# now using Wooldridge method to address it.

# ## Dynamic model with serial correlation and Wooldridge term

# We introduce in the specification of the constant a term that captures the fact that somebody is smoking during the
# first observation period.
coef_first_year = Beta('coef_first_year', 0, None, None, 0)
cte = cte_mean + coef_first_year * Smoking45 + cte_std * Draws('agent', 'NORMAL_ANTI')


# 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_wooldridge = estimate(
    log_probability, 'dynamic_model_serial_wooldridge_truncated_t_55', database
)
display(r_wooldridge)


# ### Comparison of the estimates
summary = pd.concat(
    [true_parameters['Value'], r_serial_dynamic['Value'], r_wooldridge['Value']], axis=1
)
summary.columns = [
    'True',
    'Dynamic + serial',
    'Wooldridge',
]
summary.fillna('')
display(summary)


# The estimates of the coefficients `coef_price` and `beta_last_year` are now closer to their true value. We perform
# a $t$-test analysis, to test the hypothesis that the value of the parameter is equal to its true value.
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_wooldridge.loc[param, 'Value'] - true_value) / r_wooldridge.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}')


# Except for `cte_std`, the $t$-test are, in absolute value, below 1.96. It means that we cannot reject the null
# hypothesis that the value of the parameter is equal to the true value, at the 95% level of confidence. It is
# important to realize from the relatively large value of the standard errors that the precision of the estimates is
# not high.   This is due to a lack of observations. Indeed, we observe each individual only during 10 years.
# But the Wooldridge correction has allowed to address the endogeneity issue.