"""
File: 02-outlier.py

Michel Bierlaire
Sat Aug 02 2025, 18:18:01
"""

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


# Consider as outliers all observations such that the choice probability is less than 10%.
#
# - Investigate the data in order to understand why the model is performing so poorly on those observations.
# - Use your conclusions to improve the model specification.