"""
File: 04-asv_solution.py

Michel Bierlaire
Tue Aug 05 2025, 15:57:43
"""

import biogeme.biogeme_logging as blog
from IPython.core.display_functions import display
from biogeme import models
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta, Draws, Expression, MonteCarlo, log
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)

# The variables of the model are available from the file `airline_variables.py`.
from airline_variables import (
    Fare_1,
    Fare_2,
    Fare_3,
    Legroom_1,
    Legroom_2,
    Legroom_3,
    Opt1_SchedDelayEarly,
    Opt1_SchedDelayLate,
    Opt2_SchedDelayEarly,
    Opt2_SchedDelayLate,
    Opt3_SchedDelayEarly,
    Opt3_SchedDelayLate,
    TripTimeHours_1,
    TripTimeHours_2,
    TripTimeHours_3,
    chosenAlternative,
    database,
)

# 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)
logger.info('Example b01logit.py')


# **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 10_000, say, and set `recycle` to False (if `recycle` is set to True, and a
# pickle file with the results exists, the results are imported from the file, and the estimation is not performed.)
# Then, execute the script overnight.  </div>
number_of_draws = 10
recycle = False


# # The model

# Parameters
constant2 = Beta('constant2', 0, None, None, 0)
constant3 = Beta('constant3', 0, None, None, 0)
fare = Beta('fare', 0, None, None, 0)
legroom = Beta('legroom', 0, None, None, 0)
schedule_delay_early = Beta('schedule_delay_early', 0, None, None, 0)
schedule_delay_late = Beta('schedule_delay_late', 0, None, None, 0)
total_tt1 = Beta('total_tt1', 0, None, None, 0)
total_tt2 = Beta('total_tt2', 0, None, None, 0)
total_tt3 = Beta('total_tt3', 0, None, None, 0)


# Error components
sigma_1 = Beta('sigma_1', 1, None, None, 0)
ec_1 = sigma_1 * Draws('ec_1', 'NORMAL')
sigma_2 = Beta('sigma_2', 1, None, None, 0)
ec_2 = sigma_2 * Draws('ec_2', 'NORMAL')
sigma_3 = Beta('sigma_3', 1, None, None, 0)
ec_3 = sigma_3 * Draws('ec_3', 'NORMAL')


# Utility functions. The normalization identifies the id of the alternative such that the error component is not
# included, which is equivalent to normalize its scale to zero.
def utility(normalization: int | None) -> dict[int, Expression]:
    """Utility functions.

    :param normalization: The normalization identifies the id of the
        alternative such that the error component is not included,
        which is equivalent to normalize its scale to zero. If None, no normalization is considered.
    :return: dict of utility functions
    """
    opt1 = (
        fare * Fare_1
        + legroom * Legroom_1
        + schedule_delay_early * Opt1_SchedDelayEarly
        + schedule_delay_late * Opt1_SchedDelayLate
        + total_tt1 * TripTimeHours_1
    )
    if normalization != 1:
        opt1 += ec_1
    opt2 = (
        constant2
        + fare * Fare_2
        + legroom * Legroom_2
        + schedule_delay_early * Opt2_SchedDelayEarly
        + schedule_delay_late * Opt2_SchedDelayLate
        + total_tt2 * TripTimeHours_2
    )
    if normalization != 2:
        opt2 += ec_2
    opt3 = (
        constant3
        + fare * Fare_3
        + legroom * Legroom_3
        + schedule_delay_early * Opt3_SchedDelayEarly
        + schedule_delay_late * Opt3_SchedDelayLate
        + total_tt3 * TripTimeHours_3
    )
    if normalization != 3:
        opt3 += ec_3
    return {1: opt1, 2: opt2, 3: opt3}


# # Estimation
def estimate(
    normalization: int | None, init_betas: dict[str, float] | None
) -> EstimationResults:
    """Estimate a normalized model

    :param normalization: The normalization identifies the id of the
        alternative such that the error component is not included,
        which is equivalent to normalize its scale to zero. If None,
        the model is not normalized.
    :param init_betas: if not None, starting values for the estimation.
    :return: estimation results
    """
    prob = models.logit(utility(normalization=normalization), None, chosenAlternative)
    logprob: Expression = log(MonteCarlo(prob))
    if init_betas is not None:
        logprob.change_init_values(init_betas)
    biogeme = BIOGEME(database, logprob, number_of_draws=number_of_draws)
    if normalization is None:
        biogeme.model_name = '04-asv'
    else:
        biogeme.model_name = f'04-asv_{normalization}'
    results: EstimationResults = biogeme.estimate(recycle=recycle)
    return results


# ## Estimation without normalization
unnormalized_results: EstimationResults = estimate(normalization=None, init_betas=None)

# General statistics
print(unnormalized_results.short_summary())

# Number of draws are reported
reported_number_of_draws = unnormalized_results.number_of_draws
print(f'{reported_number_of_draws}')

# Final log likelihood
LL = unnormalized_results.final_loglikelihood
print(f'{LL:.6g}')

# Estimated parameters
betas = get_pandas_estimated_parameters(estimation_results=unnormalized_results)
display(betas)


# If the number of draws is sufficiently high, we see that $\sigma_2^2$ is lesser than $\sigma_1^2$ and $\sigma_3^3$.
# Therefore, the correct normalization consists in setting $\sigma_2 = 0$.

# We extract the sigmas
sigma1_value = unnormalized_results.get_parameter_value('sigma_1')
sigma2_value = unnormalized_results.get_parameter_value('sigma_2')
sigma3_value = unnormalized_results.get_parameter_value('sigma_3')


# Relevant quantities from the variance-covariance matrix: \begin{align}s_3 &= \sigma_1^2 + \sigma_2^2 \\ s_2 &=
# \sigma_1^2 + \sigma_3^2 \\ s_1 &= \sigma_2^2 + \sigma_3^2\end{align}
s1 = sigma2_value**2 + sigma3_value**2
s2 = sigma1_value**2 + sigma3_value**2
s3 = sigma1_value**2 + sigma2_value**2
print(f's1 = {s1:.3g}')
print(f's2 = {s2:.3g}')
print(f's3 = {s3:.3g}')


# ## Results with normalization

# Normalize $\sigma_1=0$.
normalized_results_1: EstimationResults = estimate(
    normalization=1, init_betas=unnormalized_results.get_beta_values()
)

# General statistics
print(normalized_results_1.short_summary())

# Estimated parameters
display(get_pandas_estimated_parameters(estimation_results=normalized_results_1))

# Normalize $\sigma_2=0$.
normalized_results_2: EstimationResults = estimate(
    normalization=2, init_betas=unnormalized_results.get_beta_values()
)

# General statistics
print(normalized_results_2.short_summary())

# Estimated parameters
display(get_pandas_estimated_parameters(estimation_results=normalized_results_2))


# Normalize $\sigma_3=0$.
normalized_results_3: EstimationResults = estimate(
    normalization=3, init_betas=unnormalized_results.get_beta_values()
)

# General statistics
print(normalized_results_3.short_summary())

# Estimated parameters
display(get_pandas_estimated_parameters(estimation_results=normalized_results_3))


normalized_results = [normalized_results_1, normalized_results_2, normalized_results_3]


# We extract the final log likelihood and the three scale parameters
def extract(the_results: EstimationResults) -> tuple[float, float, float, float]:
    """Extracts the final log likelihood and the three scale parameters.

    :param the_results: estimation results.
    :return: log likelihood and the three scale parameters.
    """
    the_final_ll = the_results.final_loglikelihood
    try:
        s_1 = the_results.get_parameter_value('sigma_1')
    except ValueError:
        s_1 = 0
    try:
        s_2 = the_results.get_parameter_value('sigma_2')
    except ValueError:
        s_2 = 0

    try:
        s_3 = the_results.get_parameter_value('sigma_3')
    except ValueError:
        s_3 = 0
    return the_final_ll, s_1, s_2, s_3


# We gather all the results together
all_results = [extract(r) for r in normalized_results]
for a_normalization in [1, 2, 3]:
    print(f'Results for sigma_{a_normalization} = 0')
    print(f'\tFinal log likelihood: {all_results[a_normalization-1][0]:.6g}')
    print(f'\tsigma_1: {all_results[a_normalization-1][1]:.3g}')
    print(f'\tsigma_2: {all_results[a_normalization-1][2]:.3g}')
    print(f'\tsigma_3: {all_results[a_normalization-1][3]:.3g}')


# The following function calculates, for a given normalization, the relevant quantities from the variance-covariance
# matrix: \begin{align}s_3 &= \sigma_1^2 + \sigma_2^2 \\ s_2 &= \sigma_1^2 + \sigma_3^2 \\ s_1 &=
# \sigma_2^2 + \sigma_3^2\end{align}
def relevant_quantities(normalization: int) -> tuple[float, float, float]:
    """Calculates, for a given normalization, the relevant quantities
    from the variance-covariance matrix

    :param normalization: The normalization identifies the id of the
        alternative such that the error component is not included,
        which is equivalent to normalize its scale to zero. If None,
        the model is not normalized.^
    :return: the three relevant quantities.
    """
    s_1 = (
        all_results[normalization - 1][2] ** 2 + all_results[normalization - 1][3] ** 2
    )
    s_2 = (
        all_results[normalization - 1][1] ** 2 + all_results[normalization - 1][3] ** 2
    )
    s_3 = (
        all_results[normalization - 1][1] ** 2 + all_results[normalization - 1][2] ** 2
    )
    return s_1, s_2, s_3


# We now compare the relevant quantities of the variance-covariance matrix

# Correct normalization: $\sigma_2 = 0$.
s1_normalization_2, s2_normalization_2, s3_normalization_2 = relevant_quantities(
    normalization=2
)
print(f's1 = {s1_normalization_2:.3g}')
print(f's2 = {s2_normalization_2:.3g}')
print(f's3 = {s3_normalization_2:.3g}')

# Same quantities from the unnormalized model
print(f's1 = {s1:.3g}')
print(f's2 = {s2:.3g}')
print(f's3 = {s3:.3g}')

# Difference
print(f'delta s1 = {s1_normalization_2 - s1:.3g}')
print(f'delta s2 = {s2_normalization_2 - s2:.3g}')
print(f'delta s3 = {s3_normalization_2 - s3:.3g}')

# The differences are due to simulation errors.

# Incorrect normalization: $\sigma_1 = 0$
s1_normalization_1, s2_normalization_1, s3_normalization_1 = relevant_quantities(
    normalization=1
)
print(f's1 = {s1_normalization_1:.3g}')
print(f's2 = {s2_normalization_1:.3g}')
print(f's3 = {s3_normalization_1:.3g}')


# Difference with unnormalized model
print(f'delta s1 = {s1_normalization_1 - s1:.3g}')
print(f'delta s2 = {s2_normalization_1 - s2:.3g}')
print(f'delta s3 = {s3_normalization_1 - s3:.3g}')


# The level of magnitude of the differences is higher than for the correct normalization, and cannot be attributed
# to simulation error.

# Incorrect normalization: $\sigma_3=0$
s1_normalization_3, s2_normalization_3, s3_normalization_3 = relevant_quantities(
    normalization=3
)
print(f's1 = {s1_normalization_3:.3g}')
print(f's2 = {s2_normalization_3:.3g}')
print(f's3 = {s3_normalization_3:.3g}')


# Difference with unnormalized model
print(f'delta s1 = {s1_normalization_3 - s1:.3g}')
print(f'delta s2 = {s2_normalization_3 - s2:.3g}')
print(f'delta s3 = {s3_normalization_3 - s3:.3g}')

# In this case, the differences are lower, and it is difficult to say if they are due to simulation or normalization
# error.