"""

biogeme.results_processing
==========================

Examples of use of several functions.

This is designed for programmers who need examples of use of the
functions of the module. The examples are designed to illustrate the
syntax. They do not correspond to any meaningful model.

Michel Bierlaire
Mon Oct 7 18:46:39 2024
"""

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 Beta, Variable, exp
from biogeme.results_processing import (
    EstimateVarianceCovariance,
    EstimationResults,
    calc_p_value,
    compile_estimation_results,
    get_f12,
    get_html_estimated_parameters,
    get_latex_estimated_parameters,
    get_pandas_correlation_results,
    get_pandas_estimated_parameters,
)
from biogeme.version import get_text

# %%
# Version of Biogeme.
print(get_text())

# %%
# Definition of a database
df = pd.DataFrame(
    {
        'Person': [1, 1, 1, 2, 2],
        'Exclude': [0, 0, 1, 0, 1],
        'Variable1': [1, 2, 3, 4, 5],
        'Variable2': [10, 20, 30, 40, 50],
        'Choice': [1, 2, 3, 1, 2],
        'Av1': [0, 1, 1, 1, 1],
        'Av2': [1, 1, 1, 1, 1],
        'Av3': [0, 1, 1, 1, 1],
    }
)
display(df)

# %%
my_data = Database('test', df)

# %%
# Definition of various expressions
Variable1 = Variable('Variable1')
Variable2 = Variable('Variable2')
beta1 = Beta('beta1', -1.0, -3, 3, 0)
beta2 = Beta('beta2', 2.0, -3, 10, 0)
likelihood = -(beta1**2) * Variable1 - exp(beta2 * beta1) * Variable2 - beta2**4
simul = beta1 / Variable1 + beta2 / Variable2
dict_of_expressions = {'log_like': likelihood, 'beta1': beta1, 'simul': simul}

# %%
# Creation of the BIOGEME object
my_biogeme = BIOGEME(my_data, dict_of_expressions, bootstrap_samples=100)
my_biogeme.model_name = 'simple_example'
results = my_biogeme.estimate(run_bootstrap=True)
print(results)

# %%
# Dump results on a file
the_yaml_file = my_biogeme.yaml_filename
results.dump_yaml_file(filename=the_yaml_file)
print(the_yaml_file)

# %%
# Results can be imported from a file previously generated
read_results = EstimationResults.from_yaml_file(filename=the_yaml_file)
print(read_results)

# %%
# Results can be formatted in LaTeX
print(get_latex_estimated_parameters(estimation_results=read_results))

# %%
# Results can be formatted in HTML
print(get_html_estimated_parameters(estimation_results=read_results))

# %%
# General statistics, including a suggested format.
statistics = read_results.get_general_statistics()
display(statistics)

# %%
# Estimated parameters as pandas dataframe
pandas_parameters = get_pandas_estimated_parameters(estimation_results=read_results)
display(pandas_parameters)

# %%
# Correlation results
pandas_correlation = get_pandas_correlation_results(estimation_results=read_results)
display(pandas_correlation)

# %%
# Obtain the values of the parameters
beta_values = read_results.get_beta_values()
display(beta_values)

# %%
# Obtain the value of one or several specific parameters
some_beta_values = read_results.get_beta_values(my_betas=['beta2'])
display(some_beta_values)

# %%
# Variance-covariance matrix (Rao-Cramer)
rao_cramer = read_results.get_variance_covariance_matrix(
    variance_covariance_type=EstimateVarianceCovariance.RAO_CRAMER
)
display(rao_cramer)
# %%
# Variance-covariance matrix (robust)
robust = read_results.get_variance_covariance_matrix(
    variance_covariance_type=EstimateVarianceCovariance.ROBUST
)
display(robust)

# %%
# Variance-covariance matrix (bootstrap)
bootstrap = read_results.get_variance_covariance_matrix(
    variance_covariance_type=EstimateVarianceCovariance.BOOTSTRAP
)
display(bootstrap)

# %%
# Draws for sensitivity analysis are generated using
# bootstrapping. Any indicator can be generated by the model for each
# draw, and its empirical distribution can be investigated.
read_results.get_betas_for_sensitivity_analysis(['beta1', 'beta2'], size=10)

# %%
# Results can be produced in the ALOGIT F12 format
print(get_f12(estimation_results=read_results))

# %%
# Miscellaneous functions

# %%
# Likelihood ratio test. Let's first estimate a constrained model

# %%
beta2_constrained = Beta('beta2_constrained', 2.0, -3, 10, 1)
likelihood_constrained = (
    -(beta1**2) * Variable1
    - exp(beta2_constrained * beta1) * Variable2
    - beta2_constrained**4
)
my_biogeme_constrained = BIOGEME(my_data, likelihood_constrained)
my_biogeme_constrained.model_name = 'simple_example_constrained'
results_constrained = my_biogeme_constrained.estimate()
print(results_constrained.short_summary())

# %%
# We can now perform a likelihood ratio test.
test_results = results.likelihood_ratio_test(results_constrained, 0.95)
print(test_results.message)
print(f'Statistic: {test_results.statistic}')
print(f'Threshold: {test_results.threshold}')

# %%
# Calculation of the :math:`p`-value
calc_p_value(1.96)

# %%
# Compilation of results
dict_of_results = {'Model A': read_results, 'Model B': the_yaml_file}

# %%
df = compile_estimation_results(dict_of_results)
display(df)