"""
File: 04-elasticities.py

Michel Bierlaire
Sun Aug 03 2025, 17:27:36
"""

from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.models import logit, loglogit
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)

from optima_specification import v_base
from optima_variables import (
    Choice,
    database,
)

# The objective of this laboratory is to use an estimated choice model to calculate elasticities.
logprob = loglogit(v_base, None, Choice)
model_name = 'optima_base'


# If the model has already been estimated, we read the results from file.
biogeme = BIOGEME(database, logprob)
biogeme.model_name = model_name
results: EstimationResults = biogeme.estimate(run_bootstrap=True, recycle=True)

# General statistics.
print(results.short_summary())

# Estimated parameters.
display(get_pandas_estimated_parameters(estimation_results=results))


# # Choice probabilities
prob_pt = logit(v_base, None, 0)
prob_car = logit(v_base, None, 1)
prob_sm = logit(v_base, None, 2)


# # Disaggregate elasticities

# First, calculate the following disaggregate elasticities for all individuals in the sample.

# The direct point elasticity of travel time for public transportation is defined as
# $$ E_\text{time}^{P_n(\text{PT})} = \frac{\partial P_n(\text{PT})}{\partial \text{time}}
# \frac{\text{time}}{P_n(\text{PT})}.$$

# The derivative of an expression with respect to a literal is calculated using the expression `Derive`.
# For instance, $\frac{\partial P_n(\text{PT})}{\partial \text{time}}$ is calculated as
# ```
# Derive(prob_pt, 'TimePT')
# ```

# The direct point elasticity of travel cost for public transportation is defined as
# $$ E_\text{cost}^{P_n(\text{PT})} = \frac{\partial P_n(\text{PT})}{\partial \text{cost}}
# \frac{\text{cost}}{P_n(\text{PT})}.$$

# The direct point elasticity of travel time for car is defined as
# $$ E_\text{time}^{P_n(\text{car})} = \frac{\partial P_n(\text{car})}{\partial \text{time}}
# \frac{\text{time}}{P_n(\text{car})}.$$

# The direct point elasticity of travel cost for car is defined as
# $$ E_\text{cost}^{P_n(\text{car})} = \frac{\partial P_n(\text{car})}{\partial \text{cost}}
# \frac{\text{cost}}{P_n(\text{car})}.$$

# The direct point elasticity of distance for slow modes is defined as
# $$ E_\text{distance}^{P_n(\text{SM})} = \frac{\partial P_n(\text{SM})}{\partial \text{distance}}
# \frac{\text{distance}}{P_n(\text{SM})}.$$

# Analyze more specifically the values for individual number 0 and individual number 2 in the sample.

# For individual 0:
# - What would be the probability of choosing public transportation if the travel cost of this mode increases by 1%?
# - What would be the probability of choosing public transportation if the travel time of this mode increases by 1%?
# - What would be the probability of choosing slow mode for a trip that is 1% longer?

# For individual 2: what would be the probability of choosing public transportation if the travel cost of this mode
# increases by 1%?

# Aggregate elasticities capture the relative change at the level of the market shares. They can be derived from
# disaggregate elasticities using the following formula:
# $$E^{\widehat{population_shares}(i)}_{x_{jk}}  =\frac{1}{\sum_{\ell=1}^{N} \omega_\ell P_\ell(i)}\sum_{n=1}^{N}\omega_n
#   P_n(i) E^{P_n(i)}_{x_{jnk}}.$$
#

# Calculate the aggregate elasticities for travel cost and travel time for public transportation. Can the demand be
# considered elastic with respect to these two attributes?