"""
File: 01-logit_airline.py

Michel Bierlaire
Sat Aug 02 2025, 17:19:33
"""

import pandas as pd
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.database import Database
from biogeme.expressions import Beta, Variable
from biogeme.models import boxcox, loglogit, piecewise_formula
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)

# The goal of this computer session is to estimate a logit model with more than two alternatives, and to investigate
# several specifications.

# We are using stated preference data for the choice of an airline itinerary. The data set is available as
# [transp-or.epfl.ch/data/airline.dat](http://transp-or.epfl.ch/data/airline.dat), and its description is
# available [here](http://transp-or.epfl.ch/documents/technicalReports/CS_AirlineDescription.pdf).

# We recommend to download the dataset in your local directory.
df = pd.read_csv('airline.dat', sep='\t')
display(df)

# Import the database for Biogeme.
database = Database('airline', df)

# We identify the columns that will be used as variable in our model.
ArrivalTimeHours_1 = Variable('ArrivalTimeHours_1')
ArrivalTimeHours_2 = Variable('ArrivalTimeHours_2')
ArrivalTimeHours_3 = Variable('ArrivalTimeHours_3')
BestAlternative_1 = Variable('BestAlternative_1')
BestAlternative_2 = Variable('BestAlternative_2')
BestAlternative_3 = Variable('BestAlternative_3')

q11_DepartureOrArrivalIsImportant = Variable('q11_DepartureOrArrivalIsImportant')
q12_IdealDepTime = Variable('q12_IdealDepTime')
q13_IdealArrTime = Variable('q13_IdealArrTime')
DepartureTimeMins_1 = Variable('DepartureTimeMins_1')
DepartureTimeMins_2 = Variable('DepartureTimeMins_2')
DepartureTimeMins_3 = Variable('DepartureTimeMins_3')
ArrivalTimeMins_1 = Variable('ArrivalTimeMins_1')
ArrivalTimeMins_2 = Variable('ArrivalTimeMins_2')
ArrivalTimeMins_3 = Variable('ArrivalTimeMins_3')

Fare_1 = Variable('Fare_1')
Fare_2 = Variable('Fare_2')
Fare_3 = Variable('Fare_3')
Legroom_1 = Variable('Legroom_1')
Legroom_2 = Variable('Legroom_2')
Legroom_3 = Variable('Legroom_3')
TripTimeHours_1 = Variable('TripTimeHours_1')
TripTimeHours_2 = Variable('TripTimeHours_2')
TripTimeHours_3 = Variable('TripTimeHours_3')

# We define new variables from existing ones.

# The chosen alternative
chosen_alternative = (
    (BestAlternative_1 * 1) + (BestAlternative_2 * 2) + (BestAlternative_3 * 3)
)

# Sensitivity to departure or arrival time (binary variable)
departure_time_sensitive = q11_DepartureOrArrivalIsImportant == 1
arrival_time_sensitive = q11_DepartureOrArrivalIsImportant == 2
missing_sensitive = (q11_DepartureOrArrivalIsImportant != 1) * (
    q11_DepartureOrArrivalIsImportant != 2
)

# Desired departure and arrival time
desired_departure_time = q12_IdealDepTime
desired_arrival_time = q13_IdealArrTime

# The scheduled delay measures the difference between the desired and the actual value for the departure or
# arrival time.
scheduled_delay_1 = departure_time_sensitive * (
    DepartureTimeMins_1 - desired_departure_time
) + arrival_time_sensitive * (ArrivalTimeMins_1 - desired_arrival_time)
scheduled_delay_2 = departure_time_sensitive * (
    DepartureTimeMins_2 - desired_departure_time
) + arrival_time_sensitive * (ArrivalTimeMins_2 - desired_arrival_time)
scheduled_delay_3 = departure_time_sensitive * (
    DepartureTimeMins_3 - desired_departure_time
) + arrival_time_sensitive * (ArrivalTimeMins_3 - desired_arrival_time)

# We make a distinction between being earlier and later, compared to the preferred time.
opt1_sched_delay_early = -scheduled_delay_1 * (scheduled_delay_1 < 0) / 60
opt2_sched_delay_early = -scheduled_delay_2 * (scheduled_delay_2 < 0) / 60
opt3_sched_delay_early = -scheduled_delay_3 * (scheduled_delay_3 < 0) / 60

opt1_sched_delay_late = scheduled_delay_1 * (scheduled_delay_1 > 0) / 60
opt2_sched_delay_late = scheduled_delay_2 * (scheduled_delay_2 > 0) / 60
opt3_sched_delay_late = scheduled_delay_3 * (scheduled_delay_3 > 0) / 60

# # Model specification

# We are now ready to specify the choice model. We start with a simple model, that contains only generic coefficients:
# $$\begin{align*}
# V_class_1 &= \beta_\text{fare}  \text{fare}_1 + \beta_\text{legroom}  \text{legroom}_1
#        + \beta_\text{sd\_early} \text{sched\_delay\_early}_1  + \beta_\text{sd\_late} \text{sched\_delay\_late}_1
#        + \beta_\text{time} \text{elapsed\_time}_1 \\
# V_class_2 &= \text{cte}_2 + \beta_\text{fare}  \text{fare}_2 + \beta_\text{legroom}  \text{legroom}_2
#        + \beta_\text{sd\_early} \text{sched\_delay\_early}_2  + \beta_\text{sd\_late} \text{sched\_delay\_late}_2
#        + \beta_\text{time} \text{elapsed\_time}_2 \\
# V_3 &= \text{cte}_3 + \beta_\text{fare}  \text{fare}_3 + \beta_\text{legroom}  \text{legroom}_3
#        + \beta_\text{sd\_early} \text{sched\_delay\_early}_3  + \beta_\text{sd\_late} \text{sched\_delay\_late}_3
#        + \beta_\text{time} \text{elapsed\_time}_3
# \end{align*}$$

# We define the unknown parameters using the Biogeme expression `Beta`, that takes 5 parameters:
#
# - the name of the parameter (it is advised to use the exact same name for the corresponding Python variable),
# - the starting value for the estimation (usually, 0),
# - a lower bound on the value of the coefficient, or `None` for no bound,
# - an upper bound, or `None`for no bound,
# - a parameter that is 1 if the value of the parameter must be fixed to its starting value, and 0 if it has to be
#   estimated.
constant_2 = Beta('constant_2', 0, None, None, 0)
constant_3 = Beta('constant_3', 0, None, None, 0)
beta_fare = Beta('beta_fare', 0, None, None, 0)
beta_legroom = Beta('beta_legroom', 0, None, None, 0)
beta_schedule_delay_early = Beta('beta_schedule_delay_early', 0, None, None, 0)
beta_schedule_delay_late = Beta('beta_schedule_delay_late', 0, None, None, 0)
beta_elapsed_time = Beta('beta_elapsed_time', 0, None, None, 0)

# We now write the utility functions:

# Utility functions
opt1 = (
    beta_fare * Fare_1
    + beta_legroom * Legroom_1
    + beta_schedule_delay_early * opt1_sched_delay_early
    + beta_schedule_delay_late * opt1_sched_delay_late
    + beta_elapsed_time * TripTimeHours_1
)
opt2 = (
    constant_2
    + beta_fare * Fare_2
    + beta_legroom * Legroom_2
    + beta_schedule_delay_early * opt2_sched_delay_early
    + beta_schedule_delay_late * opt2_sched_delay_late
    + beta_elapsed_time * TripTimeHours_2
)
opt3 = (
    constant_3
    + beta_fare * Fare_3
    + beta_legroom * Legroom_3
    + beta_schedule_delay_early * opt3_sched_delay_early
    + beta_schedule_delay_late * opt3_sched_delay_late
    + beta_elapsed_time * TripTimeHours_3
)

# The choice model is logit. No need to code the formula. An efficient implementation is available from Biogeme.

# We associate each utility function with the id of the corresponding alternative. It must be consistent with the
# definition of the chosen alternative.
V = {1: opt1, 2: opt2, 3: opt3}

# We obtain the choice model, that is a logit model. As Biogeme requests the contribution of each observation to the
# log likelihood function, we calculate the log of the logit model.

# The documentation of the function can be obtained with the `help`function. See also the
# [online documentation](http://biogeme.epfl.ch/sphinx/index.html).
help(loglogit)

# We obtain the formula for the log of the probability, that is the contribution of each individual to the
# log-likelihood function.
logprob = loglogit(V, None, chosen_alternative)

# We initialize Biogeme with this expression, and the database.
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'logit_airline_base'

# # Estimation of the parameters.
results: EstimationResults = biogeme.estimate()

# We first display some summary information:
print(results.short_summary())

# Then we display the estimation results
pandas_results = get_pandas_estimated_parameters(estimation_results=results)
display(pandas_results)

# The results are also available in an HTML file than can be opened in your preferred browser.
display(biogeme.html_filename)


# You now need to improve the model. We propose to investigate the following specifications.

# 1. Make the coefficients of elapsed time alternative specific.
# 2. Interact the cost coefficient with income. Not that income is not available for all observations. Therefore, we
#    propose to estimate a separate coefficient for observations where income is available. For other observations,
#    define $$\beta_\text{fare} = \beta_\text{fare\_base} / \text{Cont\_Income}.$$
# 3. Investigate nonlinear specifications for the travel time variable:
#      - Consider a piecewise linear specification. This model is available with Biogeme.
#      - Consider a Box-Cox transform.
#      - Consider a power series of degree 2.

# Note that Biogeme already implements the functions for piecewise linear specification, and Box-Cox transform.
help(piecewise_formula)

#
help(boxcox)