"""
File: 08-lecture.py

Michel Bierlaire
Thu Aug 07 2025, 08:35:26

"""

import biogeme.biogeme_logging as blog
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,
    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)

# The objective of this series of exercises is to perform a similar modeling exercise as seen during the lecture.

# # 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)

# The starting point is the logit model presented above.
#
# 1. Load the results of the model where the travel time coefficient is normally distributed within the population.
#
# 2. Load the results of the model where the travel time coefficient is log normally distributed within the population.
#
# 3. Load the results of the model with two latent classes: one where the travel time coefficient is constrained to
#    be zero, and one where the travel time coefficient is estimated.
#
# 4. Estimate the same latent class model, including a class membership model, that involves the following binary
#    variables:
#
#     - `MALE`: 1 if individual is male.
#
#     - `GA`: 1 if individual owns a yearly subscription.
#
#     - `PURPOSE == 3`: 1 if trip purpose is business.
#
#     - `INCOME <= 1`: 1 if the income is less or equal to 50kCHF per year.
#
#     - `FIRST`: 1 if the individual is a first class traveler.
#
# 5. Calculate, for each segment of the population, the probability to be in each class.
#
# 6. Compare the results