"""
File: 01-stratified_sampling.py

Michel Bierlaire
Mon Aug 04 2025, 10:27:55
"""

import numpy as np
import pandas as pd
from pandas import DataFrame, Series

# The objective of this laboratory is to experiment various sampling strategies.
#
# A synthetic population composed of 1 million individuals has been generated based on the Optima database,
# containing real observations, real individuals, with observed values of all the variables.
# Each synthetic individual was generated as follows:
#
#  - One real individual was randomly selected.
#  - The socio-economic characteristics of the synthetic individual are the same as those of the real individuals.
#  - The attributes of each alternative are calculated from the real observed attributes, modified with a random perturbation.
#  - The chosen alternative is generated by simulation using the "true" choice model specified above, with the "true" value of the parameters.
#
# The scripts that have been used are
#   - `generate_population_variables.py` for the generation of the dependent variables,
#   - `generate_population_choices.py` for the generation of the choices.

population_data_file = 'synthetic_population_with_choice.zip'
population = pd.read_csv(population_data_file)
population.describe()


# Calculate the population size, and set the sample size.
population_size = population.shape[0]
print(f'Population size: {population_size:_}')
choice_variable = 'Choice'
choice_set = np.sort(population[choice_variable].unique())
print(f'{choice_set=}')

# Sample size used for these experiments
sample_size = 1000

# We describe the procedure to select a sample from that population using a stratified sample, where the strata are
# defined based on the trip purpose, the language and the chosen alternative. Then, you will be asked to implement a
# similar procedure for other sampling strategies.

# ## Step 1: Stratification.

# For each segment, we define a function that generates a Pandas mask

# In the population, there are two trip purposes: 1: 'work', 2: 'not_work'
print(population['TripPurpose'].unique())

# There are also two linguistic regions: 1: 'french', 2: 'german'
print(population['LangCode'].unique())

# The strata are defined based on the trip purpose, the language and the chosen alternative.   As the choice is
# involved in the definition of strata, the stratified sampling is *endogenous*.

# We define first the strata based on exogenous variables (trip purpose and language)


# Trip purpose = work (1), language = French (1)
def mask_work_french(data_frame: DataFrame) -> Series:
    """Generate a mask for trip purpose = work (1), language = French (1)

    :param data_frame: data frame where the mask is applied
    """
    return (data_frame['TripPurpose'] == 1) & (data_frame['LangCode'] == 1)


# Check the number and share of individuals in the population in this segment.
# Using the notation in the lecture, the share is $W(x) = N(x) / N$
population_work_french = mask_work_french(population).sum()
population_share_work_french = population_work_french / population_size
print(f'Size segment work, French in the population:  {population_work_french}')
print(
    f'Share segment work, French in the population: W(x) = {100*population_share_work_french:.3g}%'
)


# Trip purpose = work (1), language = German (2)
def mask_work_german(data_frame: DataFrame) -> Series:
    """Generate a mask for trip purpose = work (1), language = German (2)

    :param data_frame: data frame where the mask is applied
    """
    return (data_frame['TripPurpose'] == 1) & (data_frame['LangCode'] == 2)


# Check the number of individuals in the population in this segment
population_work_german = mask_work_german(population).sum()
population_share_work_german = population_work_german / population_size
print(f'Size segment work, German in the population:  {population_work_german}')
print(
    f'Share segment work, German in the population: W(x) = {100*population_share_work_german:.3g}%'
)


# Trip purpose = non work (2), language = French (1)
def mask_non_work_french(data_frame: DataFrame) -> Series:
    """Generate a mask for trip purpose = non work (2), language = French (1)

    :param data_frame: data frame where the mask is applied
    """
    return (data_frame['TripPurpose'] == 2) & (data_frame['LangCode'] == 1)


# Check the number of individuals in the population in this segment
population_non_work_french = mask_non_work_french(population).sum()
population_share_non_work_french = population_non_work_french / population_size
print(f'Size segment non work, French in the population:  {population_non_work_french}')
print(
    f'Share segment non work, French in the population: W(x) = {100*population_share_non_work_french:.3g}%'
)


# Trip purpose = non work (2), language = German (2)
def mask_non_work_german(data_frame: DataFrame) -> Series:
    """Generate a mask for trip purpose = non work (1), language = German (2)

    :param data_frame: data frame where the mask is applied
    """
    return (data_frame['TripPurpose'] == 2) & (data_frame['LangCode'] == 2)


# Check the number of individuals in the population in this segment
population_non_work_german = mask_non_work_german(population).sum()
population_share_non_work_german = population_non_work_german / population_size
print(f'Size segment non work, German in the population:  {population_non_work_german}')
print(
    f'Share segment non work, German in the population: W(x) = {100*population_share_non_work_german:.3g}%'
)

# We can verify that it is indeed covering the whole population
total = (
    population_share_work_french
    + population_share_work_german
    + population_share_non_work_french
    + population_share_non_work_german
)
print(f'Total share: {100*total:.3g}%')

# ## Step 2: Target shares.

# We define the target shares of each segment in the sample.

# Using the notation from the lecture, those are $H(x)$.
target_work_french = 0.25
target_work_german = 0.25
target_non_work_french = 0.25
target_non_work_german = 0.25

# ## Step 3: Endogenous stratification.

# Within each exogenous segment, we stratify according to the alternative


# Work, French
def mask_work_french_choice(data_frame: DataFrame, alt_id: int) -> Series:
    """Generate a mask for trip purpose = work, language = French and choice = alt_id

    :param data_frame: data frame where the mask is applied
    """
    exogenous_mask = mask_work_french(data_frame)
    choice_mask = data_frame['Choice'] == alt_id
    return exogenous_mask & choice_mask


# Check the number of individuals in the population in this segment.
# In the notation of the course, this corresponds to $W(i, x)$
population_work_french_choice = {}
population_share_work_french_choice = {}
for alt_id in choice_set:
    population_work_french_choice[alt_id] = mask_work_french_choice(
        data_frame=population, alt_id=alt_id
    ).sum()
    population_share_work_french_choice[alt_id] = (
        population_work_french_choice[alt_id] / population_size
    )
    print(
        f'Size segment work, French, choice={alt_id} in the population:  {population_work_french_choice[alt_id]}'
    )
    print(
        f'Share segment work, French, choice={alt_id} in the population: W(i, x) = '
        f'{100*population_share_work_french_choice[alt_id]:.3g}%'
    )


# Work, German
def mask_work_german_choice(data_frame: DataFrame, alt_id: int) -> Series:
    """Generate a mask for trip purpose = work, language = German and choice = alt_id

    :param data_frame: data frame where the mask is applied
    """
    exogenous_mask = mask_work_german(data_frame)
    choice_mask = data_frame['Choice'] == alt_id
    return exogenous_mask & choice_mask


# Check the number of individuals in the population in this segment.
# In the notation of the course, this corresponds to $W(i, x)$
population_work_german_choice = {}
population_share_work_german_choice = {}
for alt_id in choice_set:
    population_work_german_choice[alt_id] = mask_work_german_choice(
        data_frame=population, alt_id=alt_id
    ).sum()
    population_share_work_german_choice[alt_id] = (
        population_work_german_choice[alt_id] / population_size
    )
    print(
        f'Size segment work, German, choice={alt_id} in the population:  {population_work_german_choice[alt_id]}'
    )
    print(
        f'Share segment work, German, choice={alt_id} in the population: W(i, x) = '
        f'{100*population_share_work_german_choice[alt_id]:.3g}%'
    )


# Non work, French
def mask_non_work_french_choice(data_frame: DataFrame, alt_id: int) -> Series:
    """Generate a mask for trip purpose = non work, language = French and choice = alt_id

    :param data_frame: data frame where the mask is applied
    """
    exogenous_mask = mask_non_work_french(data_frame)
    choice_mask = data_frame['Choice'] == alt_id
    return exogenous_mask & choice_mask


# Check the number of individuals in the population in this segment.
# In the notation of the course, this corresponds to $W(i, x)$
population_non_work_french_choice = {}
population_share_non_work_french_choice = {}
for alt_id in choice_set:
    population_non_work_french_choice[alt_id] = mask_non_work_french_choice(
        data_frame=population, alt_id=alt_id
    ).sum()
    population_share_non_work_french_choice[alt_id] = (
        population_non_work_french_choice[alt_id] / population_size
    )
    print(
        f'Size segment non work, French, choice={alt_id} in the population:  {population_non_work_french_choice[alt_id]}'
    )
    print(
        f'Share segment non_work, French, choice={alt_id} in the population: W(i, x) = '
        f'{100*population_share_non_work_french_choice[alt_id]:.3g}%'
    )


# Non work, German
def mask_non_work_german_choice(data_frame: DataFrame, alt_id: int) -> Series:
    """Generate a mask for trip purpose = non work, language = German and choice = alt_id

    :param data_frame: data frame where the mask is applied
    """
    exogenous_mask = mask_non_work_german(data_frame)
    choice_mask = data_frame['Choice'] == alt_id
    return exogenous_mask & choice_mask


# Check the number of individuals in the population in this segment.
# In the notation of the course, this corresponds to $W(i, x)$
population_non_work_german_choice = {}
population_share_non_work_german_choice = {}
for alt_id in choice_set:
    population_non_work_german_choice[alt_id] = mask_non_work_german_choice(
        data_frame=population, alt_id=alt_id
    ).sum()
    population_share_non_work_german_choice[alt_id] = (
        population_non_work_german_choice[alt_id] / population_size
    )
    print(
        f'Size segment non work, German, choice={alt_id} in the population:  {population_non_work_german_choice[alt_id]}'
    )
    print(
        f'Share segment non work, German, choice={alt_id} in the population: W(i, x) = '
        f'{100*population_share_non_work_german_choice[alt_id]:.3g}%'
    )

# ## Step 4: Target shares for each alternative

# We define the target shares of each alternative within each segment.

# In the notation used in the lecture, we provide $H(i | x)$.

target_work_french_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}
target_work_german_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}
target_non_work_french_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}
target_non_work_german_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}

# ## Step 5: Sampling probability
#
# We calculate the sampling probability for each segment, using the following formula:
# $$R(i, x) = \frac{H(i, x) N_s}{ W(i, x) N} = \frac{H(i | x) H(x) N_s}{ W(i , x) N} $$

# Work, French
sampling_work_french_choice = {}
for alt_id in choice_set:
    sampling_work_french_choice[alt_id] = (
        target_work_french_choice[alt_id]
        * target_work_french
        * sample_size
        / (population_share_work_french_choice[alt_id] * population_size)
    )
    print(
        f'Sampling probability for work, French and choice = {alt_id}: {sampling_work_french_choice[alt_id]:.3g}'
    )

# Work, German
sampling_work_german_choice = {}
for alt_id in choice_set:
    sampling_work_german_choice[alt_id] = (
        target_work_german_choice[alt_id]
        * target_work_german
        * sample_size
        / (population_share_work_german_choice[alt_id] * population_size)
    )
    print(
        f'Sampling probability for work, German and choice = {alt_id}: {sampling_work_german_choice[alt_id]:.3g}'
    )

# Non work, French
sampling_non_work_french_choice = {}
for alt_id in choice_set:
    sampling_non_work_french_choice[alt_id] = (
        target_non_work_french_choice[alt_id]
        * target_non_work_french
        * sample_size
        / (population_share_non_work_french_choice[alt_id] * population_size)
    )
    print(
        f'Sampling probability for non work, French and choice = {alt_id}: {sampling_non_work_french_choice[alt_id]:.3g}'
    )

# Non work, German
sampling_non_work_german_choice = {}
for alt_id in choice_set:
    sampling_non_work_german_choice[alt_id] = (
        target_non_work_german_choice[alt_id]
        * target_non_work_german
        * sample_size
        / (population_share_non_work_german_choice[alt_id] * population_size)
    )
    print(
        f'Sampling probability for non work, German and choice = {alt_id}: {sampling_non_work_german_choice[alt_id]:.3g}'
    )

# ## Step 6: Sampling

# To perform the sampling, we include two new columns in the data frame. The first one associates each individual with
# her sampling probability. The second generates a random draw for each individual.

column_sampling = 'sampling_proba'

for alt_id in choice_set:
    population.loc[mask_work_french_choice(population, alt_id), column_sampling] = (
        sampling_work_french_choice[alt_id]
    )
    population.loc[mask_work_german_choice(population, alt_id), column_sampling] = (
        sampling_work_german_choice[alt_id]
    )
    population.loc[mask_non_work_french_choice(population, alt_id), column_sampling] = (
        sampling_non_work_french_choice[alt_id]
    )
    population.loc[mask_non_work_german_choice(population, alt_id), column_sampling] = (
        sampling_non_work_german_choice[alt_id]
    )

column_random_draw = 'random_draw'
population[column_random_draw] = np.random.uniform(0, 1, size=population_size)

# We keep only the individuals such that the random draw is lesser than the sampling  probability.
sample = population[population[column_random_draw] < population[column_sampling]]

# ## Step 7: Analysis

# We now investigate the composition of the sample.

# Sample size
actual_sample_size = len(sample)
print(f'Sample size: {actual_sample_size}')

# We now investigate each segment defining the strategy. For each of them, we compare the actual share in the
# sample with the target share.
for alt_id in choice_set:
    the_sample_segment = mask_work_french_choice(data_frame=sample, alt_id=alt_id).sum()
    target = target_work_french * target_work_french_choice[alt_id]
    print(
        f'Work, French, alt {alt_id}: {the_sample_segment} observations '
        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'
    )

for alt_id in choice_set:
    the_sample_segment = mask_work_german_choice(data_frame=sample, alt_id=alt_id).sum()
    target = target_work_german * target_work_german_choice[alt_id]
    print(
        f'Work, German, alt {alt_id}: {the_sample_segment} observations '
        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'
    )

for alt_id in choice_set:
    the_sample_segment = mask_non_work_french_choice(
        data_frame=sample, alt_id=alt_id
    ).sum()
    target = target_non_work_french * target_non_work_french_choice[alt_id]
    print(
        f'Non work, French, alt {alt_id}: {the_sample_segment} observations '
        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'
    )

for alt_id in choice_set:
    the_sample_segment = mask_non_work_german_choice(
        data_frame=sample, alt_id=alt_id
    ).sum()
    target = target_non_work_german * target_non_work_german_choice[alt_id]
    print(
        f'Non work, German, alt {alt_id}: {the_sample_segment} observations '
        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'
    )

# # Other sampling strategies

# Now, adapt the procedure to sample using the following strategies:
#  - a simple random sample,
#  - a stratified sample, where the strata are defined based on trip purpose and car availability
#
# For each stratified sample strategy, the number of individuals in each stratum in the sample should be the same.

# Note that the variables are coded as follows:

# Trip purpose: 1: work, 2: non work
print(population['TripPurpose'].unique())

# Car availability: 1: available, 3: non available
print(population['CarAvail'].unique())