"""
File: 02-maximum_likelihood.py

Michel Bierlaire
Mon Aug 04 2025, 10:28:28
"""

import os
import pickle

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tqdm
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.database import Database
from biogeme.models import loglogit
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)
from pandas import DataFrame, Series

from spec_optima import Choice, av, v
from stratified_sampling import Population, SamplingStrategy, SegmentSample

# The objective of this laboratory is to investigate maximum likelihood estimation under various sampling strategies.

# The specification of the true model is available in `spec_optima.py`.

# The procedure to perform stratified sampling is implemented in the file `stratidied_sampling.py`.

# We implement the ingredients of the following experiment.

# One instance of the experiment consists in:
#  - sampling from the population using a given sampling strategy
#  - estimate the unknown parameters using the generated sample.
#
# This is repeated several times. The mean and the standard deviation of the estimated value of each parameter across
# all instances are then calculated. As we know the true value of each parameter (that is, the value that has been
# used to generate the synthetic data), we can perform a formal t-test to test the hypothesis of the true value
# of the parameter.

# We provide below the full experiment for the simple random sampling strategy, as well as the functions implementing
# three other sampling strategies. And we ask you to run the same
# experiment for one exogenous and two endogenous sampling strategies, and to analyze and comment the results.

# Full population
population_data_file = 'synthetic_population_with_choice.zip'
population_data = pd.read_csv(population_data_file)
population = Population(data=population_data, choice_variable='Choice')

# We set the seed so that the results are reproducible
np.random.seed(seed=90267)

# We define the sample size for all experiments
sample_size = 200

# # Sampling procedures

# ## Simple random sampling

srs_sampling = SamplingStrategy(
    name='srs', population=population, sample_size=sample_size, segments=None
)


# ## Exogenously stratified Sample

# The strata are defined based on trip purpose and car availability. As the choice is not involved in the definition
# of strata, the stratified sampling is *exogenous*. Target shares in the sample: 25% for each stratum.


# We need to create mask generators implementing those segments for Pandas.
# purposes = {1: 'work', 2: 'other'}
# car_availability = {1: 'car_available', 3: 'car_unavailable'}
def mask_work_car_available(dataframe: DataFrame):
    return (dataframe['TripPurpose'] == 1) & (dataframe['CarAvail'] == 1)


segment_work_car_available = SegmentSample(
    name=f'work_car_available',
    mask_generator=mask_work_car_available,
    target_share=0.25,
    target_share_alternatives=None,
)


def mask_work_car_unavailable(dataframe: DataFrame):
    return (dataframe['TripPurpose'] == 1) & (dataframe['CarAvail'] == 3)


segment_work_car_unavailable = SegmentSample(
    name=f'work_car_unavailable',
    mask_generator=mask_work_car_unavailable,
    target_share=0.25,
    target_share_alternatives=None,
)


def mask_other_car_available(dataframe: DataFrame):
    return (dataframe['TripPurpose'] == 2) & (dataframe['CarAvail'] == 1)


segment_other_car_available = SegmentSample(
    name=f'other_car_available',
    mask_generator=mask_other_car_available,
    target_share=0.25,
    target_share_alternatives=None,
)


def mask_other_car_unavailable(dataframe: DataFrame):
    return (dataframe['TripPurpose'] == 2) & (dataframe['CarAvail'] == 3)


segment_other_car_unavailable = SegmentSample(
    name=f'other_car_unavailable',
    mask_generator=mask_other_car_unavailable,
    target_share=0.25,
    target_share_alternatives=None,
)

xss_segments = [
    segment_work_car_available,
    segment_work_car_unavailable,
    segment_other_car_available,
    segment_other_car_unavailable,
]

xss_sampling = SamplingStrategy(
    name='xss', population=population, sample_size=sample_size, segments=xss_segments
)

print(xss_sampling)

# ## Endogenously stratified Sample 1


# The strata are defined based on car availability and the chosen alternative.  As the choice is involved in the
# definition of strata, the stratified sampling is *endogenous*. Target shares in the sample: there are 5 non empty
# strata, taking 20% each

# Note that we define the masks for the exogenous strata. And within each of them, we specify the share of each
# alternative.


def mask_car_available(dataframe: DataFrame):
    return dataframe['CarAvail'] == 1


car_available_segment = SegmentSample(
    name=f'car_available',
    mask_generator=mask_car_available,
    target_share=0.6,
    target_share_alternatives={0: 1 / 3, 1: 1 / 3, 2: 1 / 3},
)


def mask_car_unavailable(dataframe: DataFrame):
    return dataframe['CarAvail'] != 1


car_unavailable_segment = SegmentSample(
    name=f'car_unavailable',
    mask_generator=mask_car_unavailable,
    target_share=0.4,
    target_share_alternatives={0: 1 / 2, 1: 0, 2: 1 / 2},
)


ess_segments = [car_available_segment, car_unavailable_segment]


ess_sampling = SamplingStrategy(
    name='ess', population=population, sample_size=sample_size, segments=ess_segments
)
print(ess_sampling)

# ## Endogenously stratified Sample 2


# 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*. Target shares in the sample: there
# are 10 non empty strata, taking 10% each


def mask_french_car_available(dataframe: DataFrame):
    return (dataframe['LangCode'] == 1) & (dataframe['CarAvail'] == 1)


french_car_available_segment = SegmentSample(
    name=f'french_car_available',
    mask_generator=mask_french_car_available,
    target_share=0.3,
    target_share_alternatives={0: 1 / 3, 1: 1 / 3, 2: 1 / 3},
)


def mask_german_car_available(dataframe: DataFrame):
    return (dataframe['LangCode'] == 2) & (dataframe['CarAvail'] == 1)


german_car_available_segment = SegmentSample(
    name=f'german_car_available',
    mask_generator=mask_german_car_available,
    target_share=0.3,
    target_share_alternatives={0: 1 / 3, 1: 1 / 3, 2: 1 / 3},
)


def mask_french_car_unavailable(dataframe: DataFrame):
    return (dataframe['LangCode'] == 1) & (dataframe['CarAvail'] == 3)


french_car_unavailable_segment = SegmentSample(
    name=f'french_car_unavailable',
    mask_generator=mask_french_car_unavailable,
    target_share=0.2,
    target_share_alternatives={0: 1 / 2, 1: 0, 2: 1 / 2},
)


def mask_german_car_unavailable(dataframe: DataFrame):
    return (dataframe['LangCode'] == 2) & (dataframe['CarAvail'] == 3)


german_car_unavailable_segment = SegmentSample(
    name=f'german_car_unavailable',
    mask_generator=mask_german_car_unavailable,
    target_share=0.2,
    target_share_alternatives={0: 1 / 2, 1: 0, 2: 1 / 2},
)

ess2_segments = [
    french_car_available_segment,
    french_car_unavailable_segment,
    german_car_available_segment,
    german_car_unavailable_segment,
]


ess2_sampling = SamplingStrategy(
    name='ess2', population=population, sample_size=sample_size, segments=ess2_segments
)


# # Choice model

# The specification of the choice model is available in the file <code>spec_optima.py</code>. Note that it is the
# exact same specification that has been used to generate the synthetic choices.


# The following function extracts a sample from the population and estimates the choice model using this sample.
# For each sample, we estimate the parameters with ESML.
def sample_and_estimate(sampling_strategy: SamplingStrategy) -> Series:
    """Estimate the choice model

    The specification of the choice model is available in the file
    spec_optima.py. Note that it is the exact same
    specification that has been used to generate the synthetic
    choices.

    This function extracts a sample from the population and
    estimates the choice model using this sample.

    :param sampling_strategy: sampling_strategy to apply
    :return: estimated parameters
    """
    # Extract the sample using the provided function.

    the_sample = sampling_strategy.sample()
    if the_sample.sample.isna().any().any():
        nan_columns = the_sample.sample.columns[the_sample.sample.isna().any()].tolist()
        print(f'NaN values in column(s) {nan_columns}')

    # Prepare the data for Biogeme
    database = Database(the_sample.name, the_sample.sample)

    # Estimate the parameters
    logprob = loglogit(v, av, Choice)
    biogeme = BIOGEME(database, logprob, generate_html=False, generate_yaml=False)
    biogeme.model_name = the_sample.name
    results: EstimationResults = biogeme.quick_estimate()

    # Extract the estimated parameters
    pandas_results = get_pandas_estimated_parameters(estimation_results=results)
    return pandas_results['Value']


# # Code the experiment

# We first load the true values of the parameters
true_values_of_the_parameters = {
    'asc_car': -1.8610422074360626,
    'asc_car_german': -1.0188194471148966,
    'asc_pt': -1.6550877823146726,
    'asc_pt_car_unavail': 3.5259706461037417,
    'asc_pt_german': 0.21100250245252905,
    'beta_cost_car': -1.5653923336388653,
    'beta_cost_pt': -0.45824318794797375,
    'beta_dist': -1.3658487759164655,
    'beta_dist_high_school': -1.0242967065023554,
    'beta_dist_higher_education': -0.011682106903893624,
    'beta_dist_university': 0.43619104548817356,
    'beta_time': -1.6340227136243364,
    'beta_time_others': -0.588784620423224,
    'beta_time_part_time': -0.2778165351485947,
    'beta_waiting': -0.22948530603741885,
    'beta_waiting_not_work': -0.1501487521368328,
}

# We select some parameters that we will analyze more closely
selected_parameters = [
    'beta_cost_car',
    'beta_cost_pt',
    'asc_car',
    'asc_pt_car_unavail',
]


# The following procedure implements the experiment. It repeats sampling and estimation, stores all estimation results.
# When it is done, the mean and the standard deviation is calculated for each parameter, as well as the *t*-statistic.
def run_experiment(
    sampling_strategy: SamplingStrategy,
    repetitions: int,
) -> tuple[DataFrame, DataFrame]:
    """Function that implements the experiment.

    It repeats sampling and estimation, stores all estimation
    results. When it is done, the mean and the standard deviation is
    calculated for each parameter, as well as the *t*-statistic.

    :param sampling_strategy: sampling strategy to use
    :param repetitions: number of times the experiment is performed
    :return: two data frames: one with the complete estimation results
        (the estimate of each parameter for each repetition), and one
        with a comparison between the true value and the average of
        the estimated values.
    """
    file_name = f'{sampling_strategy.name}_{repetitions}.pickle'
    if os.path.exists(file_name):
        with open(file_name, 'rb') as f:
            results, comparisons = pickle.load(f)
            print(f'Loaded cached results from {file_name}')
        return results, comparisons

    list_of_results = []

    # Loop to generate rows
    for _ in tqdm.tqdm(range(repetitions)):
        row = sample_and_estimate(sampling_strategy)
        list_of_results.append(row.to_frame().T)

    # Concatenate all the collected results into a data frame
    results = pd.concat(list_of_results, ignore_index=True)

    comparisons = pd.DataFrame.from_dict(
        true_values_of_the_parameters, orient='index', columns=['True']
    )
    # Set the columns of results to match the parameter names (comparisons.index)
    results.columns = comparisons.index

    comparisons['Estimated'] = results.mean()
    comparisons['StdDev'] = results.std()
    comparisons['t-test'] = (
        comparisons['Estimated'] - comparisons['True']
    ) / comparisons['StdDev']

    with open(file_name, 'wb') as f:
        pickle.dump((results, comparisons), f)
        print(f'Saved results to {file_name}')
    return results, comparisons


# We also plot the distribution of the estimates, and compare it with the true value and the mean.
def report(results: DataFrame, comparisons: DataFrame, parameters: list[str]) -> None:
    """Prepare the plots for a selected list of parameters.

    We plot the distribution of the estimates, and compare it with the
    true value and the mean.

    :param results: results of the experiment
    :param comparisons: data frame containing the comparison between
        the true value and the empirical average of the estimates of
        each coefficient. Generated by the function
        :func:run_experiment.
    :param parameters: list of parameters for which we plot the distribution
    """
    print('*** Results ***')
    display(results)
    print('*** Comparisons ***')
    display(comparisons)
    for param in parameters:
        xrangespan = 5
        plt.figure()
        plt.title(param)
        results[param].plot.hist(
            bins=20,
            range=[
                comparisons.loc[param, 'Estimated']
                - xrangespan * comparisons.loc[param, 'StdDev'],
                comparisons.loc[param, 'Estimated']
                + xrangespan * comparisons.loc[param, 'StdDev'],
            ],
        )
        plt.axvline(
            true_values_of_the_parameters[param],
            color='r',
            linestyle='dashed',
            label='True Value',
        )
        plt.axvline(
            comparisons.loc[param, 'Estimated'],
            color='g',
            linestyle='dotted',
            label='Estimated Value',
        )
        plt.legend()


# # We have prepared all the functions. We now run the experiment

# First, we define the number of repetitions.
the_repetitions = 200

# ## Simple random sampling

# We run the experiment.
results_srs, comparisons_srs = run_experiment(
    sampling_strategy=srs_sampling,
    repetitions=the_repetitions,
)

# The complete results of the estimations are available in the following table.
display(results_srs)

# The average estimated value of the parameters is therefore:
display(results_srs.mean())

# In order to compare those values with the true value of the parameters (that is, those used to generate the data), we
# use the comparison table
display(comparisons_srs)

# The results above show the good adequacy between the mean of the estimated parameters and the true value.
# In particular, the hypothesis that the true value of the parameter is what it is cannot be rejected for any parameter.
# And when the difference is large (for instance, look at `BETA_DIST_high_school`) it corresponds to a parameter
# estimated with high standard error, that is with low precision.

# Note that the share of individuals in the population with no car available is about 5%. The parameters associated
# with this group are estimated will less precision than the others.

# Increasing the sample size and the number of repetitions would increase the precision. We have kept those values
# relatively low for the sake of computational time.


# We plot the results for selected parameters.
report(results=results_srs, comparisons=comparisons_srs, parameters=selected_parameters)

# We observe that the true value of the parameters are fairly well recovered.

# Now, run the same experiment for the following sampling strategies:
# - exogenous sampling defined by the function `xss`,
# - endogenous sampling defined by the function `ess1`,
# - endogenous sampling defined by the function `ess2`.

# Analyze the results in light of the theory seen in class .