"""
File: 01-testing_solution.py

Michel Bierlaire
Sat Aug 02 2025, 17:42:44
"""

from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta
from biogeme.models import loglogit
from biogeme.results_processing import (
    EstimationResults,
    calc_p_value,
    compile_estimation_results,
    get_pandas_correlation_results,
    get_pandas_estimated_parameters,
)
from scipy.stats import chi2, norm

from logit_airline_alt_spec import (
    Opt1 as Opt1_alt_spec,
    Opt2 as Opt2_alt_spec,
    Opt3 as Opt3_alt_spec,
    results as res_alt_spec,
)
from logit_airline_base import (
    Fare_1,
    Fare_2,
    Fare_3,
    chosenAlternative,
    database,
    results as res_base,
)
from logit_airline_boxcox import results as res_boxcox
from logit_airline_income import (
    Opt1 as Opt1_income,
    Opt2 as Opt2_income,
    Opt3 as Opt3_income,
    results as res_income,
)
from logit_airline_piecewise import (
    Opt1 as Opt1_piecewise,
    Opt2 as Opt2_piecewise,
    Opt3 as Opt3_piecewise,
    results as res_piecewise,
)

# The objective of this laboratory is to investigate various specification tests.

# Base model:
# \begin{align*}
# V_class_1 &= \beta_\text{fare}  \text{fare}_1 + \beta_\text{legroom}  \text{legroom}_1
#       + \beta_\text{sd\_early} \text{sched\_delay\_early}_1  + \beta_\text{sd\_late} \text{sched\_delay\_late}_1
#       + \beta_\text{time} \text{elapsed\_time}_1 \\
# V_class_2 &= \text{cte}_2 + \beta_\text{fare}  \text{fare}_2 + \beta_\text{legroom}  \text{legroom}_2
#       + \beta_\text{sd\_early} \text{sched\_delay\_early}_2  + \beta_\text{sd\_late} \text{sched\_delay\_late}_2
#       + \beta_\text{time} \text{elapsed\_time}_2 \\
# V_3 &= \text{cte}_3 + \beta_\text{fare}  \text{fare}_3 + \beta_\text{legroom}  \text{legroom}_3
#       + \beta_\text{sd\_early} \text{sched\_delay\_early}_3  + \beta_\text{sd\_late} \text{sched\_delay\_late}_3
#       + \beta_\text{time} \text{elapsed\_time}_3 \\
# \end{align*}
#
print(res_base.short_summary())

#
display(get_pandas_estimated_parameters(estimation_results=res_base))

# Alternative specific time coefficient:
# \begin{align*}
# V_class_1 &= \beta_\text{fare}  \text{fare}_1 + \beta_\text{legroom}  \text{legroom}_1
#         + \beta_\text{sd\_early} \text{sched\_delay\_early}_1  + \beta_\text{sd\_late} \text{sched\_delay\_late}_1
#         + \beta_\text{time, 1} \text{elapsed\_time}_1 \\
# V_class_2 &= \text{cte}_2 + \beta_\text{fare}  \text{fare}_2 + \beta_\text{legroom}  \text{legroom}_2
#         + \beta_\text{sd\_early} \text{sched\_delay\_early}_2  + \beta_\text{sd\_late} \text{sched\_delay\_late}_2
#         + \beta_\text{time, 2} \text{elapsed\_time}_2 \\
# V_3 &= \text{cte}_3 + \beta_\text{fare}  \text{fare}_3 + \beta_\text{legroom}  \text{legroom}_3
#         + \beta_\text{sd\_early} \text{sched\_delay\_early}_3  + \beta_\text{sd\_late} \text{sched\_delay\_late}_3
#         + \beta_\text{time, 2} \text{elapsed\_time}_3 \\
# \end{align*}
print(res_alt_spec.short_summary())

#
parameters_alt_spec_pandas = get_pandas_estimated_parameters(
    estimation_results=res_alt_spec
)
display(parameters_alt_spec_pandas)

#
parameters_alt_spec = res_alt_spec.get_beta_values()

# # Question 1: test the null hypothesis that `beta_elapsed_time_2` is equal to `beta_elapsed_time_3`.

# Under the null hypothesis, the statistic
# $$
# \frac{\widehat{\beta}_2 - \widehat{\beta}_3}{\sqrt{\operatorname{Var}(\widehat{\beta}_2 - \widehat{\beta}_3)}},
# $$
# is distributed as a $N(0, 1)$, where
# $$
# \operatorname{Var}(\widehat{\beta}_2 - \widehat{\beta}_3) = \operatorname{Var}(\widehat{\beta}_2) +
# \operatorname{Var}(\widehat{\beta}_3) - 2 \operatorname{Covar}(\widehat{\beta}_2, \widehat{\beta}_3).
# $$
# Note that it is called a $t$-statistic, as in theory, it follows a $t$ distribution. But with a large sample size,
# it is very well approximated by the normal distribution.

# We first calculate the numerator:
numerator = (
    parameters_alt_spec['beta_elapsed_time_2']
    - parameters_alt_spec['beta_elapsed_time_3']
)
print(f'{numerator=:.3g}')

# For the denominator, we need the covariance between the estimates. It is available in the following table.
correlation_alt_spec = get_pandas_correlation_results(estimation_results=res_alt_spec)
display(correlation_alt_spec)

#
covar_beta2_beta3 = correlation_alt_spec.loc[
    (correlation_alt_spec['First parameter'] == 'beta_elapsed_time_3')
    & (correlation_alt_spec['Second parameter'] == 'beta_elapsed_time_2'),
    'Robust covariance',
].values[0]
print(f'{covar_beta2_beta3=:.3g}')

# Therefore, the square of the denominator is calculated as follows:
std_err_beta_elapsed_time_2 = res_alt_spec.get_parameter_std_err(
    parameter_name='beta_elapsed_time_2'
)
std_err_beta_elapsed_time_3 = res_alt_spec.get_parameter_std_err(
    parameter_name='beta_elapsed_time_3'
)
squared_denominator = (
    std_err_beta_elapsed_time_2**2
    + std_err_beta_elapsed_time_3**2
    - 2 * covar_beta2_beta3
)
# And the $t$-statistics is:
t_stat = numerator / (squared_denominator**0.5)
print(f't-statistic: {t_stat:.3g}')

# Note that this value is immediately available in the table `correlation_alt_spec` generated by the estimation
# procedure. Note also that the sign of the value is irrelevant.
t_test_beta_2_beta_3 = correlation_alt_spec.loc[
    (correlation_alt_spec['First parameter'] == 'beta_elapsed_time_3')
    & (correlation_alt_spec['Second parameter'] == 'beta_elapsed_time_2'),
    'Robust t-stat.',
].values[0]
print(f'{covar_beta2_beta3=:.3g}')

print(f't-test: {t_test_beta_2_beta_3:.3g}')

# Threshold for the 10% confidence level.
print(f'{norm.ppf(0.9):.3g}')

# Given the low value of the $t$-statistic, we cannot reject the null hypothesis that `beta_elapsed_time_2` is equal
# to `beta_elapsed_time_3` at 10% confidence. Actually, as the high $p$-value suggests, the null hypothesis cannot
# be rejected here for any reasonable confidence level.
print(f'{100*calc_p_value(t_test_beta_2_beta_3):.3g}%')

# If we reject the null hypothesis, there is 94% probability to make an error.

# # Question 2: test the null hypothesis that, between the two models, the base model is the true model.

# The base model is the restricted model. It is obtained from the unrestricted model using the linear restrictions:
#     $$\beta_{t,1} = \beta_{t, 2} = \beta_{t, 3}.$$ Therefore, we can perform a likelihood ratio test.
# Under the null hypothesis, the statistic
# $$ -2(\mathcal{L}(\widehat{\beta}_R) - \mathcal{L}(\widehat{\beta}_U))$$
# is distributed as $\chi^2$ with $K_U - K_R$ degrees of freedom.

# We collect the general statistics for the restricted model.
general_statistics_base = res_base.get_general_statistics()
display(general_statistics_base)

# We collect the general statistics for the unrestricted model.
general_statistics_alt_spec = res_alt_spec.get_general_statistics()
display(general_statistics_alt_spec)

# We calculate the statistic for the test.
LR = res_base.final_loglikelihood
LU = res_alt_spec.final_loglikelihood
test = -2 * (LR - LU)
print(f'Likelihood ratio statistic: {test:.3g}')

# We calculate the number of degrees of freedom.
KR = res_base.number_of_free_parameters
KU = res_alt_spec.number_of_free_parameters
degrees_of_freedom = KU - KR
print(f'Degrees of freedom: {degrees_of_freedom}')

# What is the threshold for the statistics at a 5% significance level?
threshold = chi2.ppf(0.95, degrees_of_freedom)
print(f'Threshold for the test if 5%: {threshold:.3g}')

# As the value of the statistic is below the threshold, the null hypothesis cannot be rejected at the 5% level.

# There is a tool that allows to perform directly this analysis.
lr_result = res_alt_spec.likelihood_ratio_test(res_base, 0.05)
print(f'{lr_result.statistic=:.3g}')
print(f'{lr_result.threshold=:.3g}')
print(lr_result.message)

# # Question 3: Test the null hypothesis that the fare coefficient does not vary with income.

# Base model.
print(res_base.short_summary())

# Model where the fare coefficient varies with income.
print(res_income.short_summary())

#  Under the null hypothesis, the true model is the base model. Note that we cannot perform a likelihood ratio test
#  here, as the base model cannot be obtained from a linear restriction of the parameters of the unrestricted model.

# If we compare the results, we note that the model where the fare coefficient varies with income has a poorer fit
# than the base model, although it has one more parameter. It seems to suggest that the base model should be preferred.
# In order to formally test it, we perform a Cox test.

# Composite model: we add the term involving cost from the base model to each utility function.
beta_fare = Beta('beta_fare', 0, None, None, 0)

#
Opt1_coxtest = Opt1_income + beta_fare * Fare_1
Opt2_coxtest = Opt2_income + beta_fare * Fare_2
Opt3_coxtest = Opt3_income + beta_fare * Fare_3

#
V = {1: Opt1_coxtest, 2: Opt2_coxtest, 3: Opt3_coxtest}
logprob_coxtest = loglogit(V, None, chosenAlternative)
biogeme_coxtest = BIOGEME(database, logprob_coxtest)
biogeme_coxtest.model_name = 'logit_airline_coxtest'
results_coxtest: EstimationResults = biogeme_coxtest.estimate()

# General statistics of the composite model.
print(results_coxtest.short_summary())

# Estimated parameters of the composite model.
display(get_pandas_estimated_parameters(estimation_results=results_coxtest))

# We perform an LR test to compare the composite model (unrestricted) with the base model (restricted).
lr_composite_base = results_coxtest.likelihood_ratio_test(res_base)
print(f'{lr_composite_base.statistic=:.3g}')
print(f'{lr_composite_base.threshold=:.3g}')
print(lr_composite_base.message)

# We perform an LR test to compare the composite model (unrestricted) with the "income" model (restricted).
lr_composite_income = results_coxtest.likelihood_ratio_test(res_income)
print(f'{lr_composite_income.statistic=:.3g}')
print(f'{lr_composite_income.threshold=:.3g}')
print(lr_composite_income.message)

# Both models are rejected against the composite model. It shows that none of the two models can be accepted and that
# better models should be investigated. Note that it does not mean that the composite model is the correct one.
# This model has been specified for the sake of the test, and does not necessarily correspond to a behaviorally
# meaningful specification.

# # Question 4: test the linear specification with alternative specific time parameters, against the nonlinear specification with Box-Cox transform of the time variable.

# Alternative specific time parameters.
print(res_alt_spec.short_summary())

# Box-Cox model
print(res_boxcox.short_summary())

# There are two ways to test the null hypothesis that the true model is the linear model:
#
#  - using a $t$-test,
#  - using a likelihood ratio test.

# Under the null hypothesis, the true value of $\lambda$ is 1. Therefore, we perform a $t$-test to test that
# hypothesis. Note that the $t$-test reported in the results is to test if the true value of the parameter is zero.
# Therefore, it cannot be used here.
boxcox_estimates = get_pandas_estimated_parameters(estimation_results=res_boxcox)
display(boxcox_estimates)

# Under the null hypothesis, the statistic $$\frac{\widehat{\lambda}-1}{\widehat{\sigma}_\lambda}$$ follows
# approximately a $N(0, 1)$.
lambda_value = res_boxcox.get_parameter_value('lambda_boxcox')
lambda_std_err = res_boxcox.get_parameter_std_err('lambda_boxcox')
t_test_lambda = (lambda_value - 1) / lambda_std_err
print(f't-test testing lambda=1: {t_test_lambda:.3g}')

# The threshold for the 10% confidence level is the 95% quantile of the normal distribution:
print(f'95% quantile of the normal distribution: {norm.ppf(.95):.3g}')

# Therefore, the null hypothesis cannot be rejected at the 10% level of confidence.

# We can calculate the p-value
print(f'p-value: {calc_p_value(t_test_lambda):.3g}')

# If we reject the null hypothesis that the true model is the linear model, there is a 23.5% probability to make
# an error.

# We can also perform a likelihood ratio test.
lr_boxbox_alt_spec = res_boxcox.likelihood_ratio_test(res_alt_spec, 0.10)
print(f'{lr_boxbox_alt_spec.statistic=:.3g}')
print(f'{lr_boxbox_alt_spec.threshold=:.3g}')
print(lr_boxbox_alt_spec.message)

# We reach the same conclusion. The null hypothesis cannot be rejected at the 10% level.

#  # Test the piecewise linear specification

# First, we test the piecewise linear specification against the base model. To do so, we can perform a likelihood
# ratio test. Indeed, the base model can be obtained from the piecewise linear specification by imposing the linear
# constraints that all the time coefficients `beta_time_0_2`, `beta_time_2_3` and `beta_time_3_more` are equal.
lr_piecewise_base = res_piecewise.likelihood_ratio_test(res_base, 0.01)
print(f'{lr_piecewise_base.statistic=:.3g}')
print(f'{lr_piecewise_base.threshold=:.3g}')
print(lr_piecewise_base.message)

# It seems that the null hypothesis that the base model is the true one can be rejected at 1% level of confidence.
# Therefore, out of the two, we prefer the piecewise linear specification.

# Second, we test the piecewise linear specification against the specification with the alternative specific time
# coefficient. Both models have 9 parameters, and the log likelihood for the piecewise linear model is higher.

# We test the null hypothesis that the alternative specific model is the true model. To test this hypothesis, we need
# to perform a **Davidson McKinnon test**, where $M_1$ is the alternative specific model, and $M_2$ is the piecewise
# linear model.

# We write a composite specification.

# We fix the values of the parameters of the piecewise linear specification to their estimated value. We also change
# the name of the beta to avoid duplicates in the composite model.
piecewise_estimates = res_piecewise.get_beta_values()
print(piecewise_estimates)

#
Opt1_piecewise.fix_betas(beta_values=piecewise_estimates, prefix='piecewise_')
Opt2_piecewise.fix_betas(beta_values=piecewise_estimates, prefix='piecewise_')
Opt3_piecewise.fix_betas(beta_values=piecewise_estimates, prefix='piecewise_')

# We create the composite specification
alpha = Beta('alpha', 0, None, None, 0)
Opt1 = (1 - alpha) * Opt1_alt_spec + alpha * Opt1_piecewise
Opt2 = (1 - alpha) * Opt2_alt_spec + alpha * Opt2_piecewise
Opt3 = (1 - alpha) * Opt3_alt_spec + alpha * Opt3_piecewise

# We estimate the composite model
V = {1: Opt1, 2: Opt2, 3: Opt3}
logprob = loglogit(V, None, chosenAlternative)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'logit_airline_piecewise'
results_composite: EstimationResults = biogeme.estimate()

# Estimated parameters.
composite_parameters = get_pandas_estimated_parameters(
    estimation_results=results_composite
)
display(composite_parameters)

# $t$-test for $\alpha$.
t_test_alpha = results_composite.get_parameter_t_test('alpha')
print(f't-test for alpha: {t_test_alpha:.3g}')

# $p$-value for $\alpha$.
p_value_alpha = results_composite.get_parameter_p_value('alpha')
print(f'p-value for alpha: {p_value_alpha:.3g}')

# Under the null hypothesis, the true value of $\alpha$ is 0. This hypothesis can be safely rejected here. Therefore,
# the piecewise linear specification is preferred.

# Note that we reach the same conclusion when comparing the AIC and the BIC of both models.
comparison, _ = compile_estimation_results(
    {'Alt. spec': res_alt_spec, 'Piecewise linear': res_piecewise}
)
display(comparison)