"""
File: 05-wtp_solution.py

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

import numpy as np
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta, Derive, exp, log, logzero
from biogeme.models import loglogit
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)
from matplotlib import pyplot as plt
from pandas import DataFrame

from optima_variables import (
    Choice,
    CostCarCHF,
    MarginalCostPT,
    NbTransf,
    TimeCar,
    TimePT,
    database,
    distance_km,
    female,
    fulltime,
    male,
    normalizedWeight,
    notfulltime,
    unreportedGender,
)

# The objective of this laboratory is to calculate the willingness to pay for the improvement of some attributes.

# # First model

# In the first model, the variables involved for the calculation of the willingness to pay appear linearly in the
# utility functions.

# Parameters to be estimated.
asc_car = Beta('asc_car', 0, None, None, 0)
asc_sm = Beta('asc_sm', 0, None, None, 0)
beta_time_fulltime_pt = Beta('beta_time_fulltime_pt', 0, None, None, 0)
beta_time_other_pt = Beta('beta_time_other_pt', 0, None, None, 0)
beta_time_fulltime_car = Beta('beta_time_fulltime_car', 0, None, None, 0)
beta_time_other_car = Beta('beta_time_other_car', 0, None, None, 0)

beta_dist_male = Beta('beta_dist_male', 0, None, None, 0)
beta_dist_female = Beta('beta_dist_female', 0, None, None, 0)
beta_dist_unreported = Beta('beta_dist_unreported', 0, None, None, 0)
beta_cost = Beta('beta_cost', 0, None, None, 0)
beta_transf = Beta('beta_transf', 0, None, None, 0)

# Utility functions
v_pt = (
    beta_time_fulltime_pt * TimePT * fulltime
    + beta_time_other_pt * TimePT * notfulltime
    + beta_cost * MarginalCostPT
    + beta_transf * NbTransf
)
v_car = (
    asc_car
    + beta_time_fulltime_car * TimeCar * fulltime
    + beta_time_other_car * TimeCar * notfulltime
    + beta_cost * CostCarCHF
)
v_sm = (
    asc_sm
    + beta_dist_male * distance_km * male
    + beta_dist_female * distance_km * female
    + beta_dist_unreported * distance_km * unreportedGender
)
v = {0: v_pt, 1: v_car, 2: v_sm}

# Estimation of the parameters
logprob = loglogit(v, None, Choice)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'wtp_first_model'
results: EstimationResults = biogeme.estimate()

# General statistics
print(results.short_summary())

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

# ## Value of time

# We  calculate the value of time, that is the willingness to pay to save travel time.

# Public transportation
vot_pt = Derive(v_pt, 'TimePT') / Derive(v_pt, 'MarginalCostPT')

# Car
vot_car = Derive(v_car, 'TimeCar') / Derive(v_car, 'CostCarCHF')

# Simulation
simulate = {
    'weight': normalizedWeight,
    'WTP PT time': vot_pt,
    'WTP CAR time': vot_car,
}
biosim = BIOGEME(database, simulate)
simulated_values = biosim.simulate(results.get_beta_values())

# We first look at the value of time for the public transportation alternative.
_ = simulated_values['WTP PT time'].hist()
plt.show()

# There are two values for the willingness to pay.
print(simulated_values['WTP PT time'].unique())

# They correspond to the two segments of the population characterized by `fulltime` and `notfulltime`. As the
# specification of the utility function is linear both in `TimePT` and `MarginalCostPT`, the value of time is simply
# the ratio of the coefficients.

# For the segment `fulltime`:
numerator = results.get_parameter_value('beta_time_fulltime_pt')
denominator = results.get_parameter_value('beta_cost')
vot_pt_fulltime_min = numerator / denominator
print(f'VOT public transportation / full time: {vot_pt_fulltime_min:.3g} CHF/min.')
print(f'VOT public transportation / full time: {vot_pt_fulltime_min * 60:.3g} CHF/h.')

# According to this specification, individuals in this segment are willing to pay 0.208 CHF to save one minute of
# travel time, that is 12.5 CHF per hour.

# For the segment `notfulltime`:
numerator = results.get_parameter_value('beta_time_other_pt')
denominator = results.get_parameter_value('beta_cost')
vot_pt_not_fulltime_min = numerator / denominator
print(f'VOT public transportation / part time: {vot_pt_not_fulltime_min:.3g} CHF/min.')
print(
    f'VOT public transportation / part time: {vot_pt_not_fulltime_min * 60:.3g} CHF/h.'
)

# According to this specification, individuals in this segment are willing to pay 0.144 CHF to save one minute of
# travel time, that is 8.67 CHF per hour.

# We can do the same for the car alternative.
_ = simulated_values['WTP CAR time'].hist()
plt.show()

# There are also two values:
print(simulated_values['WTP CAR time'].unique())

# For the segment `fulltime`:
numerator = results.get_parameter_value('beta_time_fulltime_car')
denominator = results.get_parameter_value('beta_cost')
vot_car_fulltime_min = numerator / denominator
print(f'VOT car / full time: {vot_car_fulltime_min:.3g} CHF/min.')
print(f'VOT car / full time: {vot_car_fulltime_min * 60:.3g} CHF/h.')

# For the segment `notfulltime`:
numerator = results.get_parameter_value('beta_time_other_car')
denominator = results.get_parameter_value('beta_cost')
vot_car_not_fulltime_min = numerator / denominator
print(f'VOT car / part time: {vot_car_not_fulltime_min:.3g} CHF/min.')
print(f'VOT car / part time: {vot_car_not_fulltime_min * 60:.3g} CHF/h.')

# The value of time for the car alternative is 27.9 CHF/hour for individuals in the segment `fulltime` and
# 28.3 CHF/hour for individuals in the segment `nofulltime`.

# Finally, we can calculate the average value of time in the population.

# For public transportation:
avg_vot_pt = (
    simulated_values['WTP PT time'] * simulated_values['weight']
).sum() / simulated_values['weight'].sum()
print(f'Average value of time for PT: {60 * avg_vot_pt:.3g} CHF/hour')

# For car:
avg_vot_car = (
    simulated_values['WTP CAR time'] * simulated_values['weight']
).sum() / simulated_values['weight'].sum()
print(f'Average value of time for car: {60 * avg_vot_car:.3g} CHF/hour')

# ## Willingness to pay to reduce the number of transfers

# In this case, as the variable is discrete, the derivatives cannot be used. The wtp to decrease the number of
# transfers is the value $population_shares$ that solves the equation:
# $$\beta_\text{transfer} \text{NbTransf} + \beta_\text{cost} \text{MarginalCostPT} =
# \beta_\text{transfer} (\text{NbTransf}-1) + \beta_\text{cost} (\text{MarginalCostPT} + population_shares).$$

# In this case, the solution happens to be also the ratio of coefficients:
# $$population_shares=\frac{\beta_\text{transfer}}{\beta_\text{cost}}.$$
numerator = results.get_parameter_value('beta_transf')
denominator = results.get_parameter_value('beta_cost')
wtp_transfer = numerator / denominator
print(
    f'Willingness to pay to reduce the number of transfers: {wtp_transfer:.3g} CHF/transfer'
)

# # Second model: moneymetric specification

# In the moneymetric specification, the cost parameter is normalized to -1, and the scale parameter is estimated.

scale_parameter = Beta('scale_parameter', 1, 1.0e-5, None, 0)

v_pt_moneymetric = scale_parameter * (
    beta_time_fulltime_pt * TimePT * fulltime
    + beta_time_other_pt * TimePT * notfulltime
    - MarginalCostPT
    + beta_transf * NbTransf
)
v_car_moneymetric = scale_parameter * (
    asc_car
    + beta_time_fulltime_car * TimeCar * fulltime
    + beta_time_other_car * TimeCar * notfulltime
    - CostCarCHF
)
v_sm_moneymetric = scale_parameter * (
    asc_sm
    + beta_dist_male * distance_km * male
    + beta_dist_female * distance_km * female
    + beta_dist_unreported * distance_km * unreportedGender
)
v_moneymetric = {0: v_pt_moneymetric, 1: v_car_moneymetric, 2: v_sm_moneymetric}

# Estimation of the parameters
logprob = loglogit(v_moneymetric, None, Choice)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'wtp_second_model'
results_second: EstimationResults = biogeme.estimate()

# General statistics. Note that the final log likelihood is the same as for the first model. It makes sense as
# the two models are equivalent, up to a normalization.
print(results_second.short_summary())

# Estimated parameters.
parameters_second = get_pandas_estimated_parameters(estimation_results=results_second)
display(parameters_second)

# As the utility function is expressed in CHF units, the parameters provide us directly with the WTP,
# with the opposite sign. You can verify that the values are  the same as for the first model (up to possible
# rounding errors).

# Value of time, public transportation, full time.
vot_pt_full_time_min = -results_second.get_parameter_value('beta_time_fulltime_pt')
print(f'VOT public transportation / full time: {vot_pt_full_time_min:.3g} CHF/min.')
print(f'VOT public transportation / full time: {vot_pt_full_time_min * 60:.3g} CHF/h.')

# Value of time, public transportation, part time.
vot_pt_part_time_min = -results_second.get_parameter_value('beta_time_other_pt')
print(f'VOT public transportation / part time: {vot_pt_part_time_min:.3g} CHF/min.')
print(f'VOT public transportation / part time: {vot_pt_part_time_min * 60:.3g} CHF/h.')

# Value of time, car, full time.
vot_car_full_time_min = -results_second.get_parameter_value('beta_time_fulltime_car')
print(f'VOT car / full time: {vot_car_full_time_min:.3g} CHF/min.')
print(f'VOT car / full time: {vot_car_full_time_min * 60:.3g} CHF/h.')

# Value of time, car, part time.
vot_car_part_time_min = -results_second.get_parameter_value('beta_time_other_car')
print(f'VOT car / part time: {vot_car_part_time_min:.3g} CHF/min.')
print(f'VOT car / part time: {vot_car_part_time_min * 60:.3g} CHF/h.')

# Willingness to pay to reduce the transfers.
wtp_transfer = -results_second.get_parameter_value('beta_transf')
print(
    f'Willingness to pay to reduce the number of transfers: {wtp_transfer:.3g} CHF/transfer'
)

# Clearly, the aggregate values for the population can be calculated exactly like for the first model.

# # Third model: nonlinear transformation of the variables

# We use a Biogeme expression that returns the log of a variable if it is nonzero, and 0 otherwise.
# The variable is assumed to be non-negative.
v_pt = (
    beta_time_fulltime_pt * logzero(TimePT) * fulltime
    + beta_time_other_pt * logzero(TimePT) * notfulltime
    + beta_cost * logzero(MarginalCostPT)
    + beta_transf * logzero(NbTransf)
)
v_car = (
    asc_car
    + beta_time_fulltime_car * logzero(TimeCar) * fulltime
    + beta_time_other_car * logzero(TimeCar) * notfulltime
    + beta_cost * logzero(CostCarCHF)
)
v_sm = (
    asc_sm
    + beta_dist_male * distance_km * male
    + beta_dist_female * distance_km * female
    + beta_dist_unreported * distance_km * unreportedGender
)
v = {0: v_pt, 1: v_car, 2: v_sm}

# Estimation of the parameters
logprob = loglogit(v, None, Choice)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'wtp_third_model'
results_third: EstimationResults = biogeme.estimate()

# General statistics
print(results_third.short_summary())

# Estimated parameters
parameters_third = get_pandas_estimated_parameters(estimation_results=results_third)
display(parameters_third)

# ## Value of time

# We calculate the value of time, that is the willingness to pay to save travel time.
vot_pt = Derive(v_pt, 'TimePT') / Derive(v_pt, 'MarginalCostPT')
vot_car = Derive(v_car, 'TimeCar') / Derive(v_car, 'CostCarCHF')

# Simulation.
simulate = {
    'weight': normalizedWeight,
    'WTP PT time': vot_pt,
    'WTP CAR time': vot_car,
    'MarginalCostPT': MarginalCostPT,
    'CostCarCHF': CostCarCHF,
}
biosim = BIOGEME(database, simulate)
simulated_values = DataFrame(biosim.simulate(results_third.get_beta_values()))


# Note that the formula are not valid when the cost is 0. In that case, Biogeme generates `inf` or `-inf`.
# For instance:
print(simulated_values.iloc[0])

# If the cost is zero, we set the value of time to zero as well.
simulated_values['WTP PT time'] = simulated_values['WTP PT time'].where(
    simulated_values['MarginalCostPT'] != 0, 0
)
simulated_values['WTP CAR time'] = simulated_values['WTP CAR time'].where(
    simulated_values['CostCarCHF'] != 0, 0
)

# We check that the values have been set correctly.
print(simulated_values.iloc[0])

# We first look at the value of time for the public transportation alternative.
_ = simulated_values['WTP PT time'].hist(bins=200)
plt.show()

# There are now many different values for the value of time. Indeed, because of the nonlinear specification, the
# value of time depends also on the value of the variables.
print(simulated_values['WTP PT time'].value_counts())

# Consider only the nonzero values for a clearer plot.
_ = simulated_values[simulated_values['WTP PT time'] != 0]['WTP PT time'].hist(bins=200)
plt.show()

# We can do the same for the car alternative.
_ = simulated_values['WTP CAR time'].hist(bins=200)
plt.show()

# Removing the zeros...
_ = simulated_values[simulated_values['WTP CAR time'] != 0]['WTP CAR time'].hist(
    bins=200
)
plt.show()

# Finally, we can calculate the average value of time in the population.

# For public transportation:
avg_vot_pt = (
    simulated_values['WTP PT time'] * simulated_values['weight']
).sum() / simulated_values['weight'].sum()
print(f'Average value of time for PT: {60 * avg_vot_pt:.3g} CHF/hour')

# For car:
avg_vot_car = (
    simulated_values['WTP CAR time'] * simulated_values['weight']
).sum() / simulated_values['weight'].sum()
print(f'Average value of time for car: {60 * avg_vot_car:.3g} CHF/hour')

# Those values are abnormally low, triggering suspicion on the validity of the model specification.

# ## Willingness to pay to reduce the number of transfers

# Again, as the variable is discrete, the derivatives cannot be used. The wtp to decrease the number of transfers is
# the value $population_shares$ that solves the equation:
# $$\beta_\text{transfer} \log(\text{NbTransf}) + \beta_\text{cost} \log(\text{MarginalCostPT}) =
# \beta_\text{transfer} \log(\text{NbTransf}-1) + \beta_\text{cost} \log(\text{MarginalCostPT} + population_shares).$$
#
# The solution is: $$population_shares=\exp\left(\frac{\beta_\text{transfer}}{\beta_\text{cost}} [\log(\text{NbTransf}) -
# \log(\text{NbTransf}-1)] + \log(\text{MarginalCostPT})\right)-\text{MarginalCostPT}.$$
wtp_transfer = (
    exp(
        (log(NbTransf) - log(NbTransf - 1)) * beta_transf / beta_cost
        + log(MarginalCostPT)
    )
    - MarginalCostPT
)

# Simulation
simulate = {
    'weight': normalizedWeight,
    'WTP Transfers': wtp_transfer,
    'Nbr. of transfers': NbTransf,
    'Cost PT': MarginalCostPT,
}


biosim = BIOGEME(database, simulate)
simulated_values = biosim.simulate(results_third.get_beta_values())


# Note that we do not verify if the arguments of the `log` is involved in the formula are valid, that is, positive.
# When they are not, Biogeme generates either a `NaN` or `inf`.


# For example, the following entry is invalid as the number of transfers is 1, and it is not possible to take the
# log of 1-1=0.
bad_index_1 = 5
print(f'Simulated values: {simulated_values.loc[bad_index_1]}')

# The following entry is invalid as the number of transfers is 0, and it is not possible to take the log of 0-1=-1.
bad_index_0 = 1903
print(simulated_values.loc[bad_index_0])

# We set the willingness to pay to zero in these cases.
simulated_values['WTP Transfers'] = simulated_values['WTP Transfers'].replace(
    [np.nan, np.inf, -np.inf], 0
)

# The willingness to pay for problematic cases is now set to zero.
print(simulated_values.loc[bad_index_1])
print(simulated_values.loc[bad_index_0])

# We can now plot the distribution.
_ = simulated_values['WTP Transfers'].hist(bins=100)
plt.show()

# As well as the distribution of nonzero values.
_ = simulated_values[simulated_values['WTP Transfers'] > 0]['WTP Transfers'].hist(
    bins=100
)
plt.show()