"""
File: 05-lecture_solution.py

Michel Bierlaire
Tue Aug 05 2025, 16:08:02
"""

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

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

# Variables used for the specification of the Swissmetro model are defined in the file `swissmetro_variables.py`.


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


# **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 10000, say. Then, execute the script
# overnight.  </div>
number_of_draws = 10_000

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


# ## 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
b_time_s = Beta('b_time_s', 1, None, None, 0)
b_time_rnd = b_time + b_time_s * Draws('b_time_rnd', 'NORMAL')


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


# ## Model
prob = logit(v, av, CHOICE)
logprob = log(MonteCarlo(prob))


# ## Estimation
biogeme = BIOGEME(database, logprob, number_of_draws=number_of_draws)
biogeme.model_name = '02normal'
results_normal: EstimationResults = biogeme.estimate()


# ## Results

# General statistics
print(results_normal.short_summary())

# Estimated parameters
param_normal = get_pandas_estimated_parameters(estimation_results=results_normal)
display(param_normal)


# # Comparison

# We build a summary data frame.
summary = compile_estimation_results(
    {'Logit': results_logit, 'Random param. (normal)': results_normal}
)
display(summary[0])

# You may note the following:
# - the large increase of the final log likelihood with the mixture model,
# - the increase of the (average) value of time when assume the time coefficient to be distributed,
# - there is a significant portion of the normal distribution that lies in the positive values:
mean_beta_time = results_normal.get_parameter_value('b_time')
stddev_beta_time = results_normal.get_parameter_value('b_time_s')
probability_to_be_larger_than_zero = 1 - norm.cdf(
    0, loc=mean_beta_time, scale=stddev_beta_time
)
print(
    f'Probability for beta_time to be larger than 0: {100 * probability_to_be_larger_than_zero:.3g}%'
)