"""
File: 02-aggregation_solution.py

Michel Bierlaire
Sat Aug 02 2025, 18:41:06
"""

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 Variable, exp
from biogeme.results_processing import EstimationResults

from netherlands_model import (
    beta_cost,
    beta_time_rail,
    logprob,
    prob_car,
    prob_rail,
    rail_cost_euro,
    rail_time,
    v_car,
)

# The objective of this laboratory is to calculate aggregate indicators from a choice model.

# # Calculate the weights of the sample in order to perform sample enumeration.

# We consider four segments of the population, characterized by two age categories  (41 or older, and 40 or younger),
# and gender.

# The census data report the following number $N_g$ of individuals in each segment $g$.
census = {
    'male_41_more': 4092390,
    'male_40_less': 4151092,
    'female_41_more': 3984028,
    'female_40_less': 4428289,
}


# The total size $N$ of the population is the sum over all segments.
population_size = sum(census.values())
print(f'Size of the population: {population_size:_}')


# We read the data file, and consider only the revealed preferences data.
df = pd.read_csv('netherlands.dat', sep='\t')
df = df[df['sp'] == 0]
display(df)


# Identify each segment in the database.
filters = {
    'male_41_more': (df.age == 0) & (df.gender == 0),
    'male_40_less': (df.age == 1) & (df.gender == 0),
    'female_41_more': (df.age == 0) & (df.gender == 1),
    'female_40_less': (df.age == 1) & (df.gender == 1),
}


# We count the number $S_g$ of individuals in each segment $g$ in the sample.
sample_segments = {
    segment_name: segment_rows.sum() for segment_name, segment_rows in filters.items()
}
print(sample_segments)


# We check the total sample size $S$.
total_sample = sum(sample_segments.values())
print(f'Sample size: {total_sample}')


# The weight $w_g$ associated with segment $g$ is defined as
# $$
# w_g = \frac{N_g}{N}\frac{S}{S_g}.
# $$
weights = {
    segment_name: census[segment_name] * total_sample / (segment_size * population_size)
    for segment_name, segment_size in sample_segments.items()
}
print(weights)


# We insert the weight as a new column in the database.
for segment_name, segment_rows in filters.items():
    df.loc[segment_rows, 'weight'] = weights[segment_name]


# We check that the sum of the weights is the sample size.
sum_weights = df['weight'].sum()
print(f'Sum of the weights: {sum_weights}')

# Here is the database. The new column with weights is the last.
display(df)

# # Perform sample enumeration

# We obtain the model from the specification file, and estimate its parameters.
database = Database('netherlands', df)
biogeme = BIOGEME(database, logprob)
biogeme.model_name = 'netherlands_model'
results: EstimationResults = biogeme.estimate()


# We calculate the choice probability of each alternative for each observation. We include the weight in the simulation
# results.
weight = Variable('weight')
simulate = {
    'weight': weight,
    'Prob. rail': prob_rail,
    'Prob. car': prob_car,
}

# We perform the simulation
biosim = BIOGEME(database, simulate, bootstrap_samples=100)
simulated_values = biosim.simulate(results.get_beta_values())
display(simulated_values)


# Market shares are calculated using the weighted mean of the
# individual probabilities.
simulated_values['Weighted rail'] = (
    simulated_values['weight'] * simulated_values['Prob. rail']
)
simulated_values['Weighted car'] = (
    simulated_values['weight'] * simulated_values['Prob. car']
)

# For rail:
market_share_rail = simulated_values['Weighted rail'].mean()
print(f'Market share for rail: {100*market_share_rail:.1f}%')


# For car:
market_share_car = simulated_values['Weighted car'].mean()
print(f'Market share for car: {100*market_share_car:.1f}%')


# # Confidence intervals

# In order to calculate confidence intervals, we need to re-estimate the parameters using bootstrapping.
results_bootstrapping = biogeme.estimate(run_bootstrap=True)


# We obtain a sample of values for the parameters. We use them to calculate empirically (that is, using simulation),
# the 90% confidence intervals on the simulated quantities.
b = results_bootstrapping.get_betas_for_sensitivity_analysis()
left, right = biosim.confidence_intervals(b, 0.9)

# The `left`and `right` data frames have the same structure as the `simulated_values`data frame, and
# contain the lower bound (left) and the upper bound (right) of the confidence interval.
display(left)

#
display(right)

# We calculate the weighted sum for the lower bound of the interval.
left['Weighted rail'] = left['weight'] * left['Prob. rail']
left['Weighted car'] = left['weight'] * left['Prob. car']

# And for the upper bound:
right['Weighted rail'] = right['weight'] * right['Prob. rail']
right['Weighted car'] = right['weight'] * right['Prob. car']

# We calculate the lower bounds on the market shares.
left_market_share_rail = left['Weighted rail'].mean()
left_market_share_car = left['Weighted car'].mean()

# And the upper bounds.
right_market_share_rail = right['Weighted rail'].mean()
right_market_share_car = right['Weighted car'].mean()

# Market share for car
print(
    f'Market share for car: {100*market_share_car:.1f}% '
    f'CI: ['
    f'{100*left_market_share_car:.1f}%-'
    f'{100*right_market_share_car:.1f}'
    f']'
)

# Market share for rail.
print(
    f'Market share for rail: {100*market_share_rail:.1f}% '
    f'CI: ['
    f'{100*left_market_share_rail:.1f}%-'
    f'{100*right_market_share_rail:.1f}'
    f']'
)


# # Scenario: rail cost decreased by 10%

# We write the choice probability under the proposed scenario.
new_rail_cost_euro = rail_cost_euro * 0.9
v_rail_scenario = beta_cost * new_rail_cost_euro + beta_time_rail * rail_time
prob_car_scenario = 1 / (1 + exp(v_rail_scenario - v_car))
prob_rail_scenario = 1 - prob_car_scenario


# We perform the simulation, exactly as before.
simulate_scenario = {
    'weight': weight,
    'Prob. rail': prob_rail_scenario,
    'Prob. car': prob_car_scenario,
}
biosim_scenario = BIOGEME(database, simulate_scenario)
simulated_scenario = biosim_scenario.simulate(results.get_beta_values())
left_scenario, right_scenario = biosim_scenario.confidence_intervals(b, 0.9)

# We calculate the weighted probabilities.
simulated_scenario['Weighted rail'] = (
    simulated_scenario['weight'] * simulated_scenario['Prob. rail']
)
simulated_scenario['Weighted car'] = (
    simulated_scenario['weight'] * simulated_scenario['Prob. car']
)

# And the corresponding confidence intervals.
left_scenario['Weighted rail'] = left_scenario['weight'] * left_scenario['Prob. rail']
left_scenario['Weighted car'] = left_scenario['weight'] * left_scenario['Prob. car']
right_scenario['Weighted rail'] = (
    right_scenario['weight'] * right_scenario['Prob. rail']
)
right_scenario['Weighted car'] = right_scenario['weight'] * right_scenario['Prob. car']

# Finally, the market shares and the confidence intervals.
market_share_car_scenario = simulated_scenario['Weighted car'].mean()
market_share_rail_scenario = simulated_scenario['Weighted rail'].mean()
left_market_share_rail_scenario = left_scenario['Weighted rail'].mean()
left_market_share_car_scenario = left_scenario['Weighted car'].mean()
right_market_share_rail_scenario = right_scenario['Weighted rail'].mean()
right_market_share_car_scenario = right_scenario['Weighted car'].mean()

# Market share for car.
print(
    f'Market share for car (scenario): {100*market_share_car_scenario:.1f}% '
    f'CI: ['
    f'{100*left_market_share_car_scenario:.1f}%-'
    f'{100*right_market_share_car_scenario:.1f}'
    f']'
)

# Market share for rail
print(
    f'Market share for rail (scenario): {100*market_share_rail_scenario:.1f}% '
    f'CI: ['
    f'{100*left_market_share_rail_scenario:.1f}%-'
    f'{100*right_market_share_rail_scenario:.1f}'
    f']'
)

# When we decrease the price of rail by 10%, the market shares change from 62.6%/37.4% to 59.4%/40.6%.