"""File 01-contingency_table.py

Michel Bierlaire
Sat Aug 02 2025, 16:31:08


"""

# We first import various elements needed for the script.

# The Pandas package, in charge of the management of the data.
import pandas as pd

# A function to display the results in the notebook
from IPython.core.display_functions import display

# Biogeme itself.
from biogeme.biogeme import BIOGEME

# The Biogeme interface with the database.
from biogeme.database import Database

# Finally, we import some expressions needed to write the model specification.
from biogeme.expressions import Beta, Variable, log
from biogeme.results_processing import get_pandas_estimated_parameters

# The objective of this exercise is to estimate the parameters of the
# simple model introduced in the lecture, using Biogeme.

# We want to build a model that predicts the market penetration of
# electric vehicles (EV) as a function of the income level. We have a
# sample of 1000 individuals. The data is defined using the Pandas
# package.

# Each row of the database corresponds to a cell of the contingency table.
data = pd.DataFrame(
    {
        'Age': [1, 1, 2, 2, 3, 3],
        'Electric': [1, 0, 1, 0, 1, 0],
        'Total': [65, 835, 55, 1045, 5, 495],
    }
)
display(data)

# We first import the data in the Biogeme database object.
database = Database('contingency', data)

# Definition of the variables. We use the same names as the Pandas columns.
Age = Variable('Age')
Electric = Variable('Electric')
Total = Variable('Total')

# Definition of the parameters to be estimated.
pi1 = Beta('pi1', 0.5, 0, 1, 0)
pi2 = Beta('pi2', 0.5, 0, 1, 0)
pi3 = Beta('pi3', 0.5, 0, 1, 0)


# We associate with each observation the relevant parameter, depending
# on the value of the variable Age. Note that the expressions like
# `Age == 1` returns 1 if True and 0 if False. Therefore, for each row
# in the database, `pi` will be either `pi1`, `pi2` or `pi3`, depending
# on the value of `Age`.

pi = (Age == 1) * pi1 + (Age == 2) * pi2 + (Age == 3) * pi3

# The contribution of each observation to the log likelihood function
# depends on the value of the variable Electric, and must be applied
# as many times as reported by Total, that is the value of the
# corresponding cell of the contingency table.

loglike = Total * log(pi) * (Electric == 1) + Total * log(1 - pi) * (Electric == 0)

# We create an instance of Biogeme, combining the model and the data
biogeme = BIOGEME(database, loglike)
biogeme.model_name = 'contingency'

# We estimate the parameters
results = biogeme.estimate()

# We obtain the results in a pandas table
pandas_results = get_pandas_estimated_parameters(estimation_results=results)
display(pandas_results)

# We can check that it is equal to the share in eah category.
for i in (1, 2, 3):
    total_for_age = data[data['Age'] == i]['Total'].sum()
    electric_for_age = data[(data['Age'] == i) & (data['Electric'] == 1)]['Total'].sum()
    print(f'pi_{i} = {electric_for_age / total_for_age:.3g}')


# We can now predict future market shares for a scenario with a new
# distribution of age.
pi1_estimate = pandas_results.loc[pandas_results['Name'] == 'pi1', 'Value'].values[0]
pi2_estimate = pandas_results.loc[pandas_results['Name'] == 'pi2', 'Value'].values[0]
pi3_estimate = pandas_results.loc[pandas_results['Name'] == 'pi3', 'Value'].values[0]

market_share = pi1_estimate * 0.25 + pi2_estimate * 0.50 + pi3_estimate * 0.25
print(f'Predicted market share: {100*market_share:.2g}%')

# Consider now a similar data set where we have collected data per income category, coded as follows:
# 1: low, 2: medium, 3: high.
income_data = pd.DataFrame(
    {
        'Income': [1, 1, 2, 2, 3, 3],
        'Electric': [1, 0, 1, 0, 1, 0],
        'Total': [15, 200, 50, 450, 135, 150],
    }
)
display(income_data)

# 1. Estimate the parameters of the  model predicting the choice of
#   electrical vehicle as a function of income.
#
# 2. Consider a scenario where  the
#   income distribution is as follows: 7.5% of the population
#   with low income, 40% of the population with medium income
#   and 52.5% of the population with high income. Use the
#   estimated model to forecast the market share of electric
#   vehicles under this scenario.