"""
File: 08-lecture_solution.py

Michel Bierlaire
Thu Aug 07 2025, 08:37:00
"""

import itertools

import biogeme.biogeme_logging as blog
import pandas as pd
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.expressions import Beta, Expression, exp, log
from biogeme.models import logit, loglogit
from biogeme.results_processing import (
    EstimationResults,
    compile_estimation_results,
    get_pandas_estimated_parameters,
)
from pandas import Series

# Variables used for the specification of the Swissmetro model are defined in the file `swissmetro_variables.py`.
from swissmetro_variables import (
    BUSINESS,
    CAR_AV_SP,
    CAR_CO_SCALED,
    CAR_TT_SCALED,
    CHOICE,
    FIRST,
    GA,
    LOW_INC,
    MALE,
    SM_AV,
    SM_COST_SCALED,
    SM_HE_SCALED,
    SM_TT_SCALED,
    TRAIN_AV_SP,
    TRAIN_COST_SCALED,
    TRAIN_HE_SCALED,
    TRAIN_TT_SCALED,
    database,
)

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


# # Parameters
asc_car = Beta('asc_car', 0, None, None, 0)
asc_train = Beta('asc_train', 0, None, None, 0)
b_time = Beta('b_time', 0, None, None, 0)
b_cost = Beta('b_cost', 0, None, None, 0)
b_fr = Beta('b_fr', 0, None, None, 0)


# # Availability conditions
av = {1: TRAIN_AV_SP, 2: SM_AV, 3: CAR_AV_SP}


# # Logit model

# ## Utility functions
v_train = (
    asc_train
    + b_time * TRAIN_TT_SCALED
    + b_cost * TRAIN_COST_SCALED
    + b_fr * TRAIN_HE_SCALED
)
v_swissmetro = b_time * SM_TT_SCALED + b_cost * SM_COST_SCALED + b_fr * SM_HE_SCALED
v_car = asc_car + b_time * CAR_TT_SCALED + b_cost * CAR_CO_SCALED
V = {1: v_train, 2: v_swissmetro, 3: v_car}


# ## Model
logprob = loglogit(V, av, CHOICE)


# ## Estimation
biogeme = BIOGEME(database, logprob)
biogeme.model_name = '01logit'
results_logit: EstimationResults = biogeme.estimate(recycle=True)


# ## Results

# General statistics
print(results_logit.short_summary())

# Estimated parameters
param_logit = get_pandas_estimated_parameters(estimation_results=results_logit)
display(param_logit)


# # Random parameter: normal distribution

# Read the results from file
results_normal = EstimationResults.from_yaml_file(filename='02normal.yaml')


# # Random parameter: lognormal distribution

# Read the results from file
results_lognormal = EstimationResults.from_yaml_file(filename='03lognormal.yaml')


# # Latent classes

# Read the results from file
results_latent = EstimationResults.from_yaml_file(filename='04latentClass.yaml')


# # Latent classes with class membership model

# We consider again two classes in the population. The first class of individuals have considered all variables when
# making their choice. For them, the specification of the utility function is the same as for the logit model.
v_train_class_1 = (
    asc_train
    + b_time * TRAIN_TT_SCALED
    + b_cost * TRAIN_COST_SCALED
    + b_fr * TRAIN_HE_SCALED
)
v_swissmetro_class_1 = (
    b_time * SM_TT_SCALED + b_cost * SM_COST_SCALED + b_fr * SM_HE_SCALED
)
v_car_class_1 = asc_car + b_time * CAR_TT_SCALED + b_cost * CAR_CO_SCALED
v_class_1 = {1: v_train_class_1, 2: v_swissmetro_class_1, 3: v_car_class_1}


# The second class of individuals ignored the travel time variable when making the choice. Therefore, this variable is
# removed from the utility function.
v_train_class_2 = asc_train + b_cost * TRAIN_COST_SCALED + b_fr * TRAIN_HE_SCALED
v_swissmetro_class_2 = b_cost * SM_COST_SCALED + b_fr * SM_HE_SCALED
v_car_class_2 = asc_car + b_cost * CAR_CO_SCALED
v_class_2 = {1: v_train_class_2, 2: v_swissmetro_class_2, 3: v_car_class_2}


# The following parameters are involved in the class membership model.
g_intercept = Beta('g_intercept', 0, None, None, 0)
g_male = Beta('g_male', 0, None, None, 0)
g_ga = Beta('g_ga', 0, None, None, 0)
g_business = Beta('g_business', 0, None, None, 0)
g_low_inc = Beta('g_low_inc', 0, None, None, 0)
g_first = Beta('g_first', 0, None, None, 0)


# The following function returns the expressions for the class membership probabilities. If `value` is set to `True`,
# the values instead of the expressions are returned.


# Note that `w` can potentially take any real value. We have to transform it into a probability using the
# transform `1 / (1 + exp(w))`
def omega(
    male: float | Expression,
    ga: float | Expression,
    business: float | Expression,
    low_inc: float | Expression,
    first: float | Expression,
    estimates: Series | None = None,
) -> tuple[float, float] | tuple[Expression, Expression]:
    """The following function returns the expressions for the class
     membership probabilities. If `value` is set to `True`, the values
    instead of the expressions are returned.

    Note that `w` can potentially take any real value. We have to
    transform it into a probability using the transform `1 / (1 +
    exp(w))`

    :param male: value or expression of the variable MALE
    :param ga: value or expression of the variable GA
    :param business: value or expression of the variable BUSINESS
    :param low_inc: value or expression of the variable LOW_INC
    :param first: value or expression of the variable FIRST
    :param estimates: estimated value of the parameters. If None, the expression is built.
    :return: values or expressions for the class membership probability
        for each of the two classes.

    """
    # We calculate the values of the parameters. They are provided either as parameters, or are available in the
    # estimates.
    g_intercept_w = g_intercept if estimates is None else estimates['g_intercept']
    g_male_w = g_male if estimates is None else estimates['g_male']
    g_ga_w = g_ga if estimates is None else estimates['g_ga']
    g_business_w = g_business if estimates is None else estimates['g_business']
    g_low_inc_w = g_low_inc if estimates is None else estimates['g_low_inc']
    g_first_w = g_first if estimates is None else estimates['g_first']
    w = (
        g_intercept_w
        + g_male_w * male
        + g_ga_w * ga
        + g_business_w * business
        + g_low_inc_w * low_inc
        + g_first_w * first
    )
    omega_1 = 1 / (1 + exp(w))
    omega_2 = 1 - omega_1
    if estimates is None:
        return omega_1, omega_2
    return omega_1.get_value(), omega_2.get_value()


# We use the function to obtain the expressions of the class membership probabilities.
prob_class_1, prob_class_2 = omega(MALE, GA, BUSINESS, LOW_INC, FIRST)


# ## Model

# We first calculate the choice probability for each class.
choice_prob_class_1 = logit(v_class_1, av, CHOICE)
choice_prob_class_2 = logit(v_class_2, av, CHOICE)


# The choice probability is obtained by using the class membership model.
prob = prob_class_1 * choice_prob_class_1 + prob_class_2 * choice_prob_class_2
logprob = log(prob)


# ## Estimation
biogeme = BIOGEME(database, logprob)
biogeme.model_name = '05latentClass'
results_latentsocio: EstimationResults = biogeme.estimate()


# ## Results

# General statistics
print(results_latentsocio.short_summary())

# Estimated parameters
param_latentsocio = get_pandas_estimated_parameters(
    estimation_results=results_latentsocio
)
display(param_latentsocio)


# # Membership probability

# We use the `itertools.product` function to enumerate all the combinations of values of the binary variables. We
# also use the `omega` function defined above, that returns the class membership probabilities.
rows = []
variables = [' MALE', ' GA', ' BUSINESS', ' LOW_INC', ' FIRST']
estimates = param_latentsocio.set_index('Name')['Value']
for x in itertools.product([0, 1], [0, 1], [0, 1], [0, 1], [0, 1]):
    prob = omega(*x, estimates=estimates)
    prob_dict = {f'Class {i+1}': v for i, v in enumerate(prob)}
    vars_dict = dict(zip(variables, x))
    row = {**prob_dict, **vars_dict}
    rows.append(row)
simulation = pd.DataFrame(rows)
for i in range(2):
    key = f'Class {i+1}'
    simulation[key] = simulation[key].apply(lambda the_prob: f'{100 * the_prob:.1f}%')

# Here are the values
display(simulation)


# # Comparison

# We build a summary data frame.
summary = compile_estimation_results(
    {
        'Logit': results_logit,
        'Random param. (normal)': results_normal,
        'Random param. (lognormal)': results_lognormal,
        'Latent class': results_latent,
        'Latent class with class mbship': results_latentsocio,
    }
)
display(summary[0])