"""
File: 05-wtp.py

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

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

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

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


# Consider the two models introduced below. For each of them, perform the following calculations:
#
# - Calculate the value of time for each individual in the sample for the public transportation alternative, and plot the distribution.
# - Calculate the value of time for each individual in the sample for the car alternative, and plot the distribution.
# - Calculate the average value in the population of these two quantities.
# - Calculate the willingness to pay to save one transfer for each individual in the sample, and plot the distribution.

# Hints:
# - Take into account that the variable 'NbTransf' is discrete, and not continuous.
# - The non linear model involves logarithms, that require a positive argument.

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

# # Second model: linear model in willingness to pay space.

# Consider the moneymetric specification of the above model, and use it to calculate the same willingness to pay
# indicators.

# # Third model: non linear transformation of the variables

# We use a Biogeme expression that returns the log of a variable if it is non zero, 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_second_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)