"""
File: 03-wesml_solution.py

Michel Bierlaire
Tue Aug 05 2025, 11:50:56
"""

import os
import pickle
from typing import NamedTuple

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.expressions import Variable, log
from biogeme.models import loglogit
from biogeme.results_processing import 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 compare ESML and WESML for an endogenous sampling strategy.

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

# A procedure to perform stratified sampling is implemented in the file `stratified_sampling.py`.


# # Sampling strategy


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

# Size of the sample and the population
sample_size = 200

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


# We define a data structure to store the estimation results
class EstimationResults(NamedTuple):
    """Data structure to store the estimation results."""

    esml: Series
    wesml: Series
    esml_corrected: Series


# ## Endogenously stratified sample

# 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*.
# Note that the segment where the car is available, and the car is chosen, is empty.
# Target shares in the sample: there are 5 non empty strata, taking 20% each


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: 0.2, 1: 0.2, 2: 0.6},
)


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: 0.7, 1: 0, 2: 0.3},
)


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)

# # Choice model

# The specification of the choice model is available in the file <code>spec_optima</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 standard ESML,
def sample_and_estimate(
    sampling_strategy: SamplingStrategy,
) -> EstimationResults:
    """Estimate the choice model with ESML, WESML, and ESML corrected.

    # The "corrected ESML" includes the correction terms: log(R_n) = -log(W_n), where R_n is the sampling
    # probability, and W_n is its inverse, the weight (the same as used in WESML).

    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 sampling_strategy: the sampling strategy
    :return: estimated parameters

    """
    the_sample = sampling_strategy.sample()
    database = Database(the_sample.name, the_sample.sample)
    logprob = loglogit(v, av, Choice)
    biogeme_esml = BIOGEME(database, logprob, generate_html=False, generate_yaml=False)
    biogeme_esml.model_name = f'{the_sample.name}_esml'
    results_esml: EstimationResults = biogeme_esml.quick_estimate()
    pandas_results_esml = get_pandas_estimated_parameters(
        estimation_results=results_esml
    )

    formulas = {
        'log_like': logprob,
        'weight': Variable(f'{sampling_strategy.column_weight}'),
    }
    biogeme_wesml = BIOGEME(
        database, formulas, generate_html=False, generate_yaml=False
    )
    biogeme_wesml.model_name = f'{the_sample.name}_wesml'
    results_wesml: EstimationResults = biogeme_wesml.quick_estimate()
    pandas_results_wesml = get_pandas_estimated_parameters(
        estimation_results=results_wesml
    )

    v_corrected = {
        alt: utility
        + log(Variable(f'{sampling_strategy.columns_correction_factor_prefix}_{alt}'))
        for alt, utility in v.items()
    }
    logprob_corrected = loglogit(v_corrected, av, Choice)
    biogeme_esml_corrected = BIOGEME(
        database, logprob_corrected, generate_html=False, generate_yaml=False
    )
    biogeme_esml_corrected.model_name = f'{the_sample.name}_esml_corrected'
    results_esml_corrected: EstimationResults = biogeme_esml_corrected.quick_estimate()
    pandas_results_esml_corrected = get_pandas_estimated_parameters(
        estimation_results=results_esml_corrected
    )

    return EstimationResults(
        esml=pandas_results_esml['Value'],
        wesml=pandas_results_wesml['Value'],
        esml_corrected=pandas_results_esml_corrected['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,
}


# 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, 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: the sampling strategy
    :param repetitions: number of times the experiment is performed

    :return: four data frames: one with the complete estimation results
        (the estimate of each parameter for each repetition) with ESML, WESML, and ESML corrected, and one
        with a comparison between the true value and the average of
        the estimated values.

    """
    file_name = f'_03_{sampling_strategy.name}_{repetitions}.pickle'
    if os.path.exists(file_name):
        with open(file_name, 'rb') as f:
            results_esml, results_wesml, results_esml_corrected, comparisons = (
                pickle.load(f)
            )
            print(f'Loaded cached results from {file_name}')
        return results_esml, results_wesml, results_esml_corrected, comparisons

    list_of_results_esml = []
    list_of_results_wesml = []
    list_of_results_esml_corrected = []

    for _ in tqdm.tqdm(range(repetitions)):
        estimation_results = sample_and_estimate(sampling_strategy)
        row_esml = estimation_results.esml
        list_of_results_esml.append(row_esml.to_frame().T)

        row_esml_corrected = estimation_results.esml_corrected
        list_of_results_esml_corrected.append(row_esml_corrected.to_frame().T)

        row_wesml = estimation_results.wesml
        list_of_results_wesml.append(row_wesml.to_frame().T)

    results_esml = pd.concat(list_of_results_esml, ignore_index=True)
    results_esml_corrected = pd.concat(
        list_of_results_esml_corrected, ignore_index=True
    )
    results_wesml = pd.concat(list_of_results_wesml, 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_esml.columns = comparisons.index
    results_esml_corrected.columns = comparisons.index
    results_wesml.columns = comparisons.index

    comparisons['Estimated_esml'] = results_esml.mean()
    comparisons['StdDev_esml'] = results_esml.std()
    comparisons['t-test_esml'] = (
        comparisons['Estimated_esml'] - comparisons['True']
    ) / comparisons['StdDev_esml']

    comparisons['Estimated_wesml'] = results_wesml.mean()
    comparisons['StdDev_wesml'] = results_wesml.std()
    comparisons['t-test_wesml'] = (
        comparisons['Estimated_wesml'] - comparisons['True']
    ) / comparisons['StdDev_wesml']

    comparisons['Estimated_esml_corrected'] = results_esml_corrected.mean()
    comparisons['StdDev_esml_corrected'] = results_esml_corrected.std()
    comparisons['t-test_esml_corrected'] = (
        comparisons['Estimated_esml_corrected'] - comparisons['True']
    ) / comparisons['StdDev_esml_corrected']

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


# # Run the experiment
the_repetitions = 200


# ## Endogenous sampling
results_ess1_esml, results_ess1_wesml, results_ess1_esml_corrected, comparisons_ess1 = (
    run_experiment(sampling_strategy=ess_sampling, repetitions=the_repetitions)
)

# Here are the comparisons
display(comparisons_ess1)

# Analyze the above results.
# - What general observation can be made about the three estimators?
# - Focus now on the constants `asc_car` and `ASC_PT`. What observation can be made about the three estimates?
# - Evaluate the average precision of WESML and ESML (corrected). Which one is better?

# When performing the test of the hypothesis that the estimated value is equal to the true value, we observe the
# following:

# From the column `t-test_esml`, corresponding to ESML, it is seen that all parameters are correctly recovered,
# except the constants:
asc_car_true = comparisons_ess1.loc['asc_car', 'True']
asc_car_esml = comparisons_ess1.loc['asc_car', 'Estimated_esml']
asc_car_esml_ttest = comparisons_ess1.loc['asc_car', 't-test_esml']
print(f'asc_car true value: {asc_car_true:.3g}')
print(f'asc_car ESML estimate: {asc_car_esml:.3g}')
print(f'asc_car ESML estimate (t-test): {asc_car_esml_ttest:.3g}')

asc_pt_true = comparisons_ess1.loc['asc_pt', 'True']
asc_pt_esml = comparisons_ess1.loc['asc_pt', 'Estimated_esml']
asc_pt_esml_ttest = comparisons_ess1.loc['asc_pt', 't-test_esml']
print(f'asc_pt true value: {asc_pt_true:.3g}')
print(f'asc_pt ESML estimate: {asc_pt_esml:.3g}')
print(f'asc_pt ESML estimate (t-test): {asc_pt_esml_ttest:.3g}')


# From the column `t-test_wesml`, corresponding to WESML, it is seen that all parameters are correctly recovered,
# including the constants.
asc_car_true = comparisons_ess1.loc['asc_car', 'True']
asc_car_wesml = comparisons_ess1.loc['asc_car', 'Estimated_wesml']
asc_car_wesml_ttest = comparisons_ess1.loc['asc_car', 't-test_wesml']
print(f'asc_car true value: {asc_car_true:.3g}')
print(f'asc_car WESML estimate: {asc_car_wesml:.3g}')
print(f'asc_car WESML estimate (t-test): {asc_car_wesml_ttest:.3g}')

asc_pt_true = comparisons_ess1.loc['asc_pt', 'True']
asc_pt_wesml = comparisons_ess1.loc['asc_pt', 'Estimated_wesml']
asc_pt_wesml_ttest = comparisons_ess1.loc['asc_pt', 't-test_wesml']
print(f'asc_pt true value: {asc_pt_true:.3g}')
print(f'asc_pt WESML estimate: {asc_pt_wesml:.3g}')
print(f'asc_pt WESML estimate (t-test): {asc_pt_wesml_ttest:.3g}')

# From the column `t-test_esml_corrected`, corresponding to ESML with the correction of the constants, it is seen
# that all parameters are correctly recovered, including the (corrected) constants.
asc_car_true = comparisons_ess1.loc['asc_car', 'True']
asc_car_esml_corrected = comparisons_ess1.loc['asc_car', 'Estimated_esml_corrected']
asc_car_esml__corrected_ttest = comparisons_ess1.loc['asc_car', 't-test_esml_corrected']
print(f'asc_car true value: {asc_car_true:.3g}')
print(f'asc_car ESML estimate (corrected): {asc_car_esml_corrected:.3g}')
print(f'asc_car ESML estimate (corrected, t-test): {asc_car_esml__corrected_ttest:.3g}')

asc_pt_true = comparisons_ess1.loc['asc_pt', 'True']
asc_pt_esml_corrected = comparisons_ess1.loc['asc_pt', 'Estimated_esml_corrected']
asc_pt_esml_ttest_corrected = comparisons_ess1.loc['asc_pt', 't-test_esml_corrected']
print(f'asc_pt true value: {asc_pt_true:.3g}')
print(f'asc_pt ESML estimate (corrected): {asc_pt_esml_corrected:.3g}')
print(f'asc_pt ESML estimate (corrected, t-test): {asc_pt_esml_ttest_corrected:.3g}')


# If we now compare the standard deviation for the two estimators, it appears that using ESML and correcting the
# constants provided more precise estimates than WESML.
std_dev_wesml = comparisons_ess1['StdDev_wesml'].abs().mean()
std_dev_esml_corrected = comparisons_ess1['StdDev_esml_corrected'].abs().mean()
print(f'Std. dev. WESML: {std_dev_wesml:.3g}')
print(f'Std. dev. ESML (corrected): {std_dev_esml_corrected:.3g}')