"""
File: 01-nested_solution.py

Michel Bierlaire
Mon Aug 04 2025, 09:55:09
"""

from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta
from biogeme.models import loglogit, lognested
from biogeme.nests import NestsForNestedLogit, OneNestForNestedLogit
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)
from scipy.stats import chi2

from spec_swissmetro import av, v
from variables_swissmetro import CHOICE, database

# The objective of this laboratory is to compare a nested logit specification with a logit specification.


# ## Estimation of a logit model
logprob = loglogit(v, av, CHOICE)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'logit'
logit_results: EstimationResults = biogeme.estimate(recycle=True)

# General statistics
print(logit_results.short_summary())

# Estimated parameters
logit_parameters = get_pandas_estimated_parameters(estimation_results=logit_results)
display(logit_parameters)

# ## Nested logit

# There are three possibilities to partition the choice set:
#
# - [Car, Train]  and [Swissmetro],
# - [Train, Swissmetro] and [Car],
# - [Car, Swissmetro] and [Train].
#
# The first one groups existing alternatives together. The second one groups public transportation modes together.
# The third one being less intuitive, we select the two first specifications.
#

# ### Nested logit: existing alternatives

# There are two nests: existing and future. The second one contains only one alternative. Therefore, its nest parameter
# is not identified and set to 1.
mu_existing = Beta('mu_existing', 1, 0, None, 0)
existing = OneNestForNestedLogit(
    nest_param=mu_existing, list_of_alternatives=[1, 3], name='existing'
)
future = OneNestForNestedLogit(nest_param=1.0, list_of_alternatives=[2], name='future')
nests = NestsForNestedLogit(choice_set=list(v), tuple_of_nests=(existing, future))

# Estimation of the parameters
logprob = lognested(v, av, nests, CHOICE)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'nested_existing'
nested_existing_results: EstimationResults = biogeme.estimate(recycle=True)

# General statistics
print(nested_existing_results.short_summary())

# Estimated parameters
nested_existing_table = get_pandas_estimated_parameters(
    estimation_results=nested_existing_results
)
display(nested_existing_table)

# The nested parameter is greater than one, consistently with the theory.
value_mu_existing = nested_existing_results.get_parameter_value('mu_existing')
print(f'Nest parameter: {value_mu_existing:.3g}')


# If we test the null hypothesis that the true value of the nest parameter is 1, we use a $t$-test.
# Remember that the t-test provided in the estimation report are against 0, not 1, and cannot be used for this purpose.
mu_stderr = nested_existing_results.get_parameter_std_err('mu_existing')
tested_value = 1
t_test = (tested_value - value_mu_existing) / mu_stderr
print(f't_test against 1: {t_test:.3g}')


# Therefore, we can reject the null hypothesis at the 5% level. It means that we reject logit.

# We can also test the null hypothesis that the two models are equivalent using a likelihood ratio test:

# First, we calculate the statistic.
LL_logit = logit_results.final_loglikelihood
LL_nested_existing = nested_existing_results.final_loglikelihood
LR = -2 * (LL_logit - LL_nested_existing)
print(f'Statistic for the LR test: {LR:.3g}')


# Number of degrees of freedom:
dof = (
    nested_existing_results.number_of_free_parameters
    - logit_results.number_of_free_parameters
)
print(f'Degrees of freedom: {dof}')


# The threshold value of the $\chi$-square test with one degree of freedom at 5% level is:
print(f'Threshold for the test: {chi2.isf(0.05, dof):.3g}')

# The same test can also be performed directly as follows.
lr_test_existing = nested_existing_results.likelihood_ratio_test(
    logit_results, significance_level=0.05
)
print(f'{lr_test_existing.statistic=:.3g}')
print(f'{lr_test_existing.threshold=:.3g}')
print(lr_test_existing.message)


# Therefore, the null hypothesis can be rejected, and the nested logit model is preferred.

# ### Nested logit: public transportation modes


mu_public = Beta('mu_public', 1, 0, None, 0)

public = OneNestForNestedLogit(
    nest_param=mu_public, list_of_alternatives=[1, 2], name='public'
)
private = OneNestForNestedLogit(
    nest_param=1.0, list_of_alternatives=[3], name='private'
)
nests = NestsForNestedLogit(choice_set=list(v), tuple_of_nests=(public, private))

# Estimation of the parameters
logprob = lognested(v, av, nests, CHOICE)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'nested_public'
nested_public_results: EstimationResults = biogeme.estimate(recycle=True)

# General statistics
print(nested_public_results.short_summary())

# Estimated parameters
nested_public_table = get_pandas_estimated_parameters(
    estimation_results=nested_public_results
)
display(nested_public_table)


# The nest parameter is less than 1. This is inconsistent with the theory. The model is rejected.
value_mu_public = nested_public_results.get_parameter_value('mu_public')
print(f'Nest parameter: {value_mu_public:.3g}')


# In conclusion, among the three models, the nested logit model where the existing alternatives are in the same nest
# is preferred.