"""
File: 02-outlier_solution.py

Michel Bierlaire
Sat Aug 02 2025, 18:19:59
"""

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 logit, loglogit
from biogeme.results_processing import (
    EstimationResults,
    compile_estimation_results,
    get_pandas_estimated_parameters,
)

# The objective of this laboratory is to illustrate the outlier analysis. We use the Optima case study, for
# transportation mode choice in Switzerland.

# Read the data
df = pd.read_csv('optima.dat', sep='\t')
display(df)

# Prepare the data for Biogeme
database = Database('optima', df)

# Identification of the relevant variables.
Choice = Variable('Choice')
Weight = Variable('Weight')
age = Variable('age')
Gender = Variable('Gender')
TimePT = Variable('TimePT')
TimeCar = Variable('TimeCar')
TripPurpose = Variable('TripPurpose')
MarginalCostPT = Variable('MarginalCostPT')
CostCarCHF = Variable('CostCarCHF')
distance_km = Variable('distance_km')
Education = Variable('Education')
LangCode = Variable('LangCode')
OccupStat = Variable('OccupStat')
NbTransf = Variable('NbTransf')
Income = Variable('Income')
WaitingTimePT = Variable('WaitingTimePT')
CarAvail = Variable('CarAvail')
Subscription = Variable('Subscription')
GenAbST = Variable('GenAbST')
OwnHouse = Variable('OwnHouse')
NbBicy = Variable('NbBicy')

# Removing some incorrectly coded data.
exclude = (
    (age == -1)
    + (Choice == -1)
    + (CostCarCHF < 0)
    + (Income == -1)
    + (CarAvail == 3) * (Choice == 1)
) > 0
database.remove(exclude)


# We first estimate a simple model:
# \begin{align*}
# V_\text{PT} &= \text{Cte}_\text{PT} + \beta_{t, \text{PT}} \text{Time}_\text{PT}
#                 + \beta_{c, \text{PT}} \text{Cost}_\text{PT}, \\
# V_\text{Car} &= \text{Cte}_\text{Car} + \beta_{t, \text{Car}} \text{Time}_\text{Car}
#                 + \beta_{c, \text{Car}} \text{Cost}_\text{Car}, \\
# V_\text{SM} &= \beta_d \text{distance}.
# \end{align*}

# Parameters to be estimated.
asc_pt = Beta('asc_pt', 0, None, None, 0)
asc_car = Beta('asc_car', 0, None, None, 0)
beta_time_pt = Beta('beta_time_pt', 0.0, None, None, 0)
beta_time_car = Beta('beta_time_car', 0.0, None, None, 0)
beta_cost = Beta('beta_cost', 0, None, None, 0)
beta_distance = Beta('beta_distance', 0, None, None, 0)

# Utility functions
v_pt = asc_pt + beta_time_pt * TimePT + beta_cost * MarginalCostPT
v_car = asc_car + beta_time_car * TimeCar + beta_cost * CostCarCHF
v_sm = beta_distance * distance_km

v = {0: v_pt, 1: v_car, 2: v_sm}

# Availability conditions.
av = {0: 1, 1: CarAvail != 3, 2: 1}

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

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

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


# We now simulate the estimated model to obtain the contribution of each observation to the likelihood function. It is
# the probability predicted by the model to choose the actually chosen alternative.
prob_chosen = logit(v, av, Choice)

# We define a dictionary with the formulas to be simulated. Here, there is only one.
simulate = {'Prob. chosen': prob_chosen}

# We perform the simulation.
biosim = BIOGEME(database, simulate)
betas_values = results.get_beta_values()
sim_results = biosim.simulate(the_beta_values=betas_values)
display(sim_results)


# We now merge the result of the simulation with the original database.
# We also sort the rows based on the 'Prob. chosen'
merged_table = pd.concat([sim_results, database.dataframe], axis='columns').sort_values(
    by=['Prob. chosen']
)
display(merged_table)


# As we are interested in the outliers, we drop from the table all observations such that the probability predicted by
# the model is above 10%.
merged_table.drop(merged_table[merged_table['Prob. chosen'] >= 0.1].index, inplace=True)
display(merged_table)


# Let's investigate only the variables in the model.
variables = [
    'Choice',
    'TimePT',
    'MarginalCostPT',
    'TimeCar',
    'CostCarCHF',
    'distance_km',
]
display(merged_table[variables])


# We realize that alternative 2 (slow modes) has been selected for trips with a long distance. In order to
# investigate why, we extract all observations such that the alternative "slow modes" was chosen.
outliers_sm = merged_table[merged_table['Choice'] == 2]


# We now compare the distribution of the variable `NbBicy` in this table, with the distribution of the same variable
# in the full sample.

# Frequency of NbBicy in the outliers' database.
print(outliers_sm['NbBicy'].value_counts(normalize=True).sort_index().cumsum())

# Frequency of NbBicy in the complete database.
print(database.dataframe['NbBicy'].value_counts(normalize=True).sort_index().cumsum())

# We realize that, among the outliers that have chosen "slow modes", the proportion of individuals with 5 bicycles or
# more is very high, compared with the proportion in the full sample. It seems that the ownership of bicycles is a
# sign of an attraction to use slow modes, even for long distance.

# We investigate another variable: `age`.
print(outliers_sm['age'].value_counts(normalize=True).sort_index().cumsum())


# We realize that most outliers (~70%) are 45 or more. We can assume that individuals in that age range like to practice
# slow modes to keep in shape.

# In order to check these hypotheses, let's introduce the variables `NbBicy` and `age` in the model.
beta_bicy = Beta('beta_bicy', 0, None, None, 0)
beta_45_or_more = Beta('beta_45_or_more', 0, None, None, 0)

# New specification of the utility function for slow modes.
v_sm_revised = (
    beta_distance * distance_km + beta_bicy * NbBicy + beta_45_or_more * (age >= 45)
)
v_revised = {0: v_pt, 1: v_car, 2: v_sm_revised}

# Estimation of the revised model
logprob_revised = loglogit(v_revised, av, Choice)
biogeme_revised = BIOGEME(database, logprob_revised)
biogeme_revised.model_name = 'logit_optima_revised'
results_revised: EstimationResults = biogeme_revised.estimate()

# General statistics
print(results_revised.short_summary())

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

# We compare the two models
comparison, _ = compile_estimation_results(
    {'Base': results, 'Bicycles': results_revised}
)
display(comparison)

# We observe a significant improvement in the fit.