"""
File: 04-asv.py

Michel Bierlaire
Tue Aug 05 2025, 15:55:33
"""

import biogeme.biogeme_logging as blog
from IPython.core.display_functions import display
from biogeme import models
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta, Draws
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)

# The variables of the model are available from the file `airline_variables.py`.
from airline_variables import (
    Fare_1,
    Fare_2,
    Fare_3,
    Legroom_1,
    Legroom_2,
    Legroom_3,
    Opt1_SchedDelayEarly,
    Opt1_SchedDelayLate,
    Opt2_SchedDelayEarly,
    Opt2_SchedDelayLate,
    Opt3_SchedDelayEarly,
    Opt3_SchedDelayLate,
    TripTimeHours_1,
    TripTimeHours_2,
    TripTimeHours_3,
    chosenAlternative,
    database,
)

# The objective of this laboratory is to investigate various normalizations of an alternative specific variance model.

# Consider the logit model presented below.
#
# 1. Include error components to obtain a model with alternative specific variances.
# 2. Estimate the model without any normalization.
# 3. Identify the scale parameter that must be normalized to zero.
# 4. Compare the entries of the variance-covariance matrix for both models.
# 5. Perform the same analysis when the wrong scale parameter is normalized to zero.

# As the estimation time may be long, we ask Biogeme to report the details of the iterations.
logger = blog.get_screen_logger(level=blog.INFO)


# # The model

# Parameters
constant2 = Beta('constant2', 0, None, None, 0)
constant3 = Beta('constant3', 0, None, None, 0)
fare = Beta('fare', 0, None, None, 0)
legroom = Beta('legroom', 0, None, None, 0)
schedule_delay_early = Beta('schedule_delay_early', 0, None, None, 0)
schedule_delay_late = Beta('schedule_delay_late', 0, None, None, 0)
total_tt1 = Beta('total_tt1', 0, None, None, 0)
total_tt2 = Beta('total_tt2', 0, None, None, 0)
total_tt3 = Beta('total_tt3', 0, None, None, 0)


# Error components
sigma_1 = Beta('sigma_1', 1, None, None, 0)
ec_1 = sigma_1 * Draws('ec_1', 'NORMAL')
sigma_2 = Beta('sigma_2', 1, None, None, 0)
ec_2 = sigma_2 * Draws('ec_2', 'NORMAL')
sigma_3 = Beta('sigma_3', 1, None, None, 0)
ec_3 = sigma_3 * Draws('ec_3', 'NORMAL')


# Utility functions.
opt1 = (
    fare * Fare_1
    + legroom * Legroom_1
    + schedule_delay_early * Opt1_SchedDelayEarly
    + schedule_delay_late * Opt1_SchedDelayLate
    + total_tt1 * TripTimeHours_1
)
opt2 = (
    constant2
    + fare * Fare_2
    + legroom * Legroom_2
    + schedule_delay_early * Opt2_SchedDelayEarly
    + schedule_delay_late * Opt2_SchedDelayLate
    + total_tt2 * TripTimeHours_2
)
opt3 = (
    constant3
    + fare * Fare_3
    + legroom * Legroom_3
    + schedule_delay_early * Opt3_SchedDelayEarly
    + schedule_delay_late * Opt3_SchedDelayLate
    + total_tt3 * TripTimeHours_3
)
v = {1: opt1, 2: opt2, 3: opt3}


# # Estimation of the logit model
logprob = models.loglogit(v, None, chosenAlternative)
biogeme = BIOGEME(database, logprob)
biogeme.modelName = 'logit'

results: EstimationResults = biogeme.estimate()

# General statistics
print(results.short_summary())

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

# Hints:

# - It is advised to start working with a low number of draws, until the script is working well.
#   Then, increase the number of draws to 10000, say. Then, execute the script overnight.
# - In order to avoid repeating the same code, we suggest to write a Python function that takes as argument the
#   id of the alternative that is normalized:
#   `def estimate(normalization: int) -> bioResults:`
#   Then, depending on the normalization, you provide the corresponding specification to Biogeme.