"""

Simulation of a choice model
============================

We use an estimated model to perform various simulations.

:author: Michel Bierlaire, EPFL
:date: Wed Apr 12 21:04:33 2023

"""
import sys
import time
import pandas as pd
from biogeme import models
import biogeme.biogeme as bio
import biogeme.exceptions as excep
import biogeme.results as res
from optima_data import database, normalized_weight
from scenarios import scenario

# %%
# Obtain the specification for the default scenario.
# The definition of the scenarios is available in :ref:`scenarios`.
V, nests, _, _ = scenario()

V_PT = V[0]
V_CAR = V[1]
V_SM = V[2]

# %%
# Obtain the expression for the choice probability of each alternative.
prob_PT = models.nested(V, None, nests, 0)
prob_CAR = models.nested(V, None, nests, 1)
prob_SM = models.nested(V, None, nests, 2)

# Read the estimation results from the file
try:
    results = res.bioResults(pickleFile='saved_results/b02estimation.pickle')
except excep.BiogemeError:
    sys.exit(
        'Run first the script b02simulation.py '
        'in order to generate the '
        'file b02estimation.pickle.'
    )

# %%
# We now simulate various expressions on the database, and store the
# results in a Pandas dataframe.

# %%
start_time = time.time()
simulate_formulas = {
    'weight': normalized_weight.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
    'Utility PT': V_PT.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
    'Utility car': V_CAR.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
    'Utility SM': V_SM.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
    'Prob. PT': prob_PT.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
    'Prob. car': prob_CAR.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
    'Prob. SM': prob_SM.getValue_c(
        betas=results.getBetaValues(), database=database, prepareIds=True
    ),
}

# %%
simulated_values = pd.DataFrame.from_dict(
    simulate_formulas,
)

# %%
print(
    f'--- Execution time with getValue_c: '
    f'{time.time() - start_time:.2f} seconds ---'
)

# %%
# We now perform the same simulation using Biogeme. The results are
# identical, but the execution time is faster. Indeed, Biogeme
# recycles calculations performed for one expression for the other
# expressions.

# %%
# A dictionary with the requested expression must be provided to Biogeme
simulate = {
    'weight': normalized_weight,
    'Utility PT': V_PT,
    'Utility car': V_CAR,
    'Utility SM': V_SM,
    'Prob. PT': prob_PT,
    'Prob. car': prob_CAR,
    'Prob. SM': prob_SM,
}

# %%
start_time = time.time()
the_biogeme = bio.BIOGEME(database, simulate)
biogeme_simulation = the_biogeme.simulate(results.getBetaValues())

# %%
print(
    f'--- Execution time with Biogeme:    '
    f'{time.time() - start_time:.2f} seconds ---'
)

# %%
# Let's print the two results, to show that they are identical

# %%
# Without Biogeme
print(simulated_values)

# %%
# With Biogeme
print(biogeme_simulation)
