"""
File: 07-lecture_solution.py

Michel Bierlaire
Wed Aug 06 2025, 12:12:22
"""

import biogeme.biogeme_logging as blog
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta, log
from biogeme.models import logit, loglogit
from biogeme.results_processing import (
    EstimationResults,
    compile_estimation_results,
    get_pandas_estimated_parameters,
)

# Variables used for the specification of the Swissmetro model are defined in the file `swissmetro_variables.py`.
from swissmetro_variables import (
    CAR_AV_SP,
    CAR_CO_SCALED,
    CAR_TT_SCALED,
    CHOICE,
    SM_AV,
    SM_COST_SCALED,
    SM_HE_SCALED,
    SM_TT_SCALED,
    TRAIN_AV_SP,
    TRAIN_COST_SCALED,
    TRAIN_HE_SCALED,
    TRAIN_TT_SCALED,
    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)

# # Parameters
asc_car = Beta('asc_car', 0, None, None, 0)
asc_train = Beta('asc_train', 0, None, None, 0)
b_time = Beta('b_time', 0, None, None, 0)
b_cost = Beta('b_cost', 0, None, None, 0)
b_fr = Beta('b_fr', 0, None, None, 0)


# # Availability conditions
av = {1: TRAIN_AV_SP, 2: SM_AV, 3: CAR_AV_SP}


# # Logit model

# ## Utility functions
v_train = (
    asc_train
    + b_time * TRAIN_TT_SCALED
    + b_cost * TRAIN_COST_SCALED
    + b_fr * TRAIN_HE_SCALED
)
v_swissmetro = b_time * SM_TT_SCALED + b_cost * SM_COST_SCALED + b_fr * SM_HE_SCALED
v_car = asc_car + b_time * CAR_TT_SCALED + b_cost * CAR_CO_SCALED
v = {1: v_train, 2: v_swissmetro, 3: v_car}


# ## Model
logprob = loglogit(v, av, CHOICE)


# ## Estimation
biogeme = BIOGEME(database, logprob)
biogeme.model_name = '01logit'
results_logit: EstimationResults = biogeme.estimate(recycle=True)


# ## Results

# General statistics
print(results_logit.short_summary())

# Estimated parameters
param_logit = get_pandas_estimated_parameters(estimation_results=results_logit)
display(param_logit)


# # Random parameter: normal distribution

# Read the results from file
results_normal = EstimationResults.from_yaml_file(filename='02normal.yaml')


# # Random parameter: lognormal distribution

# Read the results from file
results_lognormal = EstimationResults.from_yaml_file(filename='03lognormal.yaml')


# # Latent classes

# We consider two classes in the population. The first class of individuals have considered all variables when making
# their choice. For them, the specification of the utility function is the same as for the logit model.
v_train_class_1 = (
    asc_train
    + b_time * TRAIN_TT_SCALED
    + b_cost * TRAIN_COST_SCALED
    + b_fr * TRAIN_HE_SCALED
)
v_swissmetro_class_1 = (
    b_time * SM_TT_SCALED + b_cost * SM_COST_SCALED + b_fr * SM_HE_SCALED
)
v_car_class_1 = asc_car + b_time * CAR_TT_SCALED + b_cost * CAR_CO_SCALED
v_class_1 = {1: v_train_class_1, 2: v_swissmetro_class_1, 3: v_car_class_1}


# The second class of individuals ignored the travel time variable when making the choice. Therefore, this variable
# is removed from the utility function.
v1_class_2 = asc_train + b_cost * TRAIN_COST_SCALED + b_fr * TRAIN_HE_SCALED
v2_class_2 = b_cost * SM_COST_SCALED + b_fr * SM_HE_SCALED
v3_class_2 = asc_car + b_cost * CAR_CO_SCALED
v_class_2 = {1: v1_class_2, 2: v2_class_2, 3: v3_class_2}


# The following parameter captures the probability to belong to class 1.
omega = Beta('omega', 0.5, 0, 1, 0)
prob_class_1 = omega
prob_class_2 = 1 - omega


# ## Model

# We first calculate the choice probability for each class.
choice_prob_class_1 = logit(v_class_1, av, CHOICE)
choice_prob_class_2 = logit(v_class_2, av, CHOICE)


# The choice probability is obtained by using the class membership model.
choice_prob = prob_class_1 * choice_prob_class_1 + prob_class_2 * choice_prob_class_2
logprob = log(choice_prob)


# ## Estimation
biogeme = BIOGEME(database, logprob)
biogeme.model_name = '04latentClass'
results_latent: EstimationResults = biogeme.estimate()


# ## Results

# General statistics

print(results_latent.short_summary())

# Estimated parameters

param_latent = get_pandas_estimated_parameters(estimation_results=results_latent)
display(param_latent)


# # Comparison

# We build a summary data frame.

summary = compile_estimation_results(
    {
        'Logit': results_logit,
        'Random param. (normal)': results_normal,
        'Random param. (lognormal)': results_lognormal,
        'Latent class': results_latent,
    }
)
display(summary[0])

# We note that the probability to belong to class one is...
estimated_proba_class_1 = results_latent.get_parameter_value('omega')
print(f'Probability to belong to class 1: {100 * estimated_proba_class_1:.3g}%.')