"""

Numerical integration
=====================

Calculation of a mixtures of logit models where the integral is
calculated using numerical integration.

Michel Bierlaire, EPFL
Sat Jun 28 2025, 21:09:58
"""

from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import IntegrateNormal, RandomVariable
from biogeme.models import logit

from swissmetro_one import (
    CAR_AV_SP,
    CAR_CO_SCALED,
    CAR_TT_SCALED,
    CHOICE,
    SM_AV,
    SM_COST_SCALED,
    SM_TT_SCALED,
    TRAIN_AV_SP,
    TRAIN_COST_SCALED,
    TRAIN_TT_SCALED,
    database,
)

# %%
# Parameters
asc_car = 0.137
asc_train = -0.402
asc_sm = 0
b_time = -2.26
b_time_s = 1.66
b_cost = -1.29

# %%
# Define a random parameter, normally distributed,
# designed to be used for integration
omega = RandomVariable('omega')
b_time_rnd = b_time + b_time_s * omega

# %%
# Definition of the utility functions
v_train = asc_train + b_time_rnd * TRAIN_TT_SCALED + b_cost * TRAIN_COST_SCALED
v_swissmetro = asc_sm + b_time_rnd * SM_TT_SCALED + b_cost * SM_COST_SCALED
v_car = asc_car + b_time_rnd * CAR_TT_SCALED + b_cost * CAR_CO_SCALED

# %%
# Associate utility functions with the numbering of alternatives
util = {1: v_train, 2: v_swissmetro, 3: v_car}

# %%
# Associate the availability conditions with the alternatives
av = {1: TRAIN_AV_SP, 2: SM_AV, 3: CAR_AV_SP}

# %%
# The choice model is a logit, with availability conditions
integrand = logit(util, av, CHOICE)
numerical_integral = IntegrateNormal(integrand, 'omega')

# %%
simulate = {'Numerical': numerical_integral}

# %%
biosim = BIOGEME(database, simulate)

# %%
results = biosim.simulate(the_beta_values={})
display(results)

# %%
print('Mixture of logit - numerical integration: ', results.iloc[0]['Numerical'])