"""
File: 04-elasticities_solution.py

Michel Bierlaire
Sun Aug 03 2025, 16:06:29
"""

import pickle

from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Derive, Expression
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,
    CostCarCHF,
    MarginalCostPT,
    TimeCar,
    TimePT,
    database,
    distance_km,
    normalizedWeight,
)

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


# # Choice probabilities
prob_PT: Expression = logit(v_base, None, 0)
prob_CAR: Expression = logit(v_base, None, 1)
prob_SM: Expression = logit(v_base, None, 2)


# # Disaggregate elasticities

# 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})}.$$
direct_elas_pt_time = Derive(prob_PT, 'TimePT') * TimePT / prob_PT


# 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})}.$$
direct_elas_pt_cost = Derive(prob_PT, 'MarginalCostPT') * MarginalCostPT / prob_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})}.$$
direct_elas_car_time = Derive(prob_CAR, 'TimeCar') * TimeCar / prob_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})}.$$
direct_elas_car_cost = Derive(prob_CAR, 'CostCarCHF') * CostCarCHF / prob_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})}.$$
direct_elas_sm_dist = Derive(prob_SM, 'distance_km') * distance_km / prob_SM

# Simulate the formulas
simulate = {
    'weight': normalizedWeight,
    'Prob. car': prob_CAR,
    'Prob. public transportation': prob_PT,
    'Prob. slow modes': prob_SM,
    'direct_elas_pt_time': direct_elas_pt_time,
    'direct_elas_pt_cost': direct_elas_pt_cost,
    'direct_elas_car_time': direct_elas_car_time,
    'direct_elas_car_cost': direct_elas_car_cost,
    'direct_elas_sm_dist': direct_elas_sm_dist,
    'Marginal cost PT': MarginalCostPT,
}

filename = 'elasticities.pickle'
try:
    with open(filename, 'rb') as f:
        simulated_values = pickle.load(f)
        print(f'Elasticities read from {filename}')
except FileNotFoundError:
    biosim = BIOGEME(database, simulate)
    simulated_values = biosim.simulate(results.get_beta_values())
    print(f'Elasticities calculated and saved in {filename}')
    with open(filename, 'wb') as f:
        pickle.dump(simulated_values, f)


# We first look at individual 0.
display(simulated_values.iloc[0])


# - This individual does not pay for public transportation. Therefore, the elasticity with respect to the cost of public transportation is 0.
# - The elasticity of travel time for public transportation is quite low. If the travel time increases by 1%, the probability to choose that mode decreases by 0.12%. It means that it will drop from 0.4750 down to 0.4745. When the value of the elasticity is less than one, it is common to say that the demand is "inelastic".

# Probability for public transportation
prob_pt = simulated_values.loc[simulated_values.index[0], 'Prob. public transportation']
print(f'Prob. public transportation: {prob_pt:.3g}')

# Direct elasticity wrt travel time.
elast_pt_time = simulated_values.loc[simulated_values.index[0], 'direct_elas_pt_time']
print(f'Direct elasticity wrt PT travel time: {elast_pt_time:.3g}')

# Probability if travel time increases by 1%.
new_prob_pt = prob_pt * (100 + elast_pt_time) / 100
print(f'Probability when travel time increases by 1%: {new_prob_pt:.3g}')


# - The elasticity of distance for slow modes is quite high. If the individual considers a trip that is 1% longer in
# distance, the probability to choose that mode decreases by 5.50%. As it is larger than 1, it is common to say that
# the demand is "elastic". However, as the initial probability to choose that mode is small, the absolute change is
# not large.

# Probability for slow modes
prob_sm = simulated_values.loc[simulated_values.index[0], 'Prob. slow modes']
print(f'Prob. slow modes: {prob_sm:.3g}')

# Elasticity wrt distance.
elast_sm_dist = simulated_values.loc[simulated_values.index[0], 'direct_elas_sm_dist']
print(f'Direct elasticity wrt distance: {elast_sm_dist:.3g}')

# Probability if distance increases.
new_prob_sm = prob_sm * (100 + elast_sm_dist) / 100
print(f'Probability when distance increases by 1%: {new_prob_sm:.3g}')


# We now look at individual 2.
display(simulated_values.iloc[2])


# - We note that the elasticity of travel cost for public transportation is larger than one (in absolute value).
# If the travel cost increases by 1%, the probability to choose public transportation decreases by 1.63%. It means that it will drop from 10.95% down to 10.77%. When the value of the elasticity is larger than one, it is common to say that the demand is "elastic".

# Probability for public transportation.
prob_pt = simulated_values.loc[simulated_values.index[2], 'Prob. public transportation']
print(f'Prob. public transportation: {prob_pt:.3g}')

# Elasticity wrt cost.
elast_pt_cost = simulated_values.loc[simulated_values.index[2], 'direct_elas_pt_cost']
print(f'Direct elasticity wrt cost: {elast_pt_cost:.3g}')

# Probability if cost increases by 1%
new_prob_pt = prob_pt * (100 + elast_pt_cost) / 100
print(f'Probability when cost increases by 1%: {new_prob_pt:.3g}')


# # Aggregate elasticities

# 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}}.$$
#

# We calculate the aggregate elasticities for travel cost and travel time for public transportation.

# Numerator for cost.
simulated_values['numerator_pt_cost'] = (
    simulated_values['weight']
    * simulated_values['Prob. public transportation']
    * simulated_values['direct_elas_pt_cost']
)

# Numerator for travel time.
simulated_values['numerator_pt_time'] = (
    simulated_values['weight']
    * simulated_values['Prob. public transportation']
    * simulated_values['direct_elas_pt_time']
)

# Denominator
simulated_values['denominator_pt'] = (
    simulated_values['weight'] * simulated_values['Prob. public transportation']
)

# Aggregate elasticity wrt cost
agg_elast_pt_cost = (
    simulated_values['numerator_pt_cost'].sum()
    / simulated_values['denominator_pt'].sum()
)
print(f'Aggregate elasticity wrt cost: {agg_elast_pt_cost:.3g}')

# Aggregate elasticity wrt travel time
agg_elast_pt_time = (
    simulated_values['numerator_pt_time'].sum()
    / simulated_values['denominator_pt'].sum()
)
print(f'Aggregate elasticity wrt travel time: {agg_elast_pt_time:.3g}')


# We note that the demand is inelastic for both those indicators, as they are less than one (in absolute value).