"""
File: 03-indicators.py

Michel Bierlaire
Sat Aug 02 2025, 18:44:59
"""

from typing import NamedTuple

import numpy as np
from IPython.core.display_functions import display
from biogeme.biogeme import BIOGEME
from biogeme.models import logit, loglogit
from biogeme.results_processing import (
    EstimationResults,
    get_pandas_estimated_parameters,
)

from optima_specification import scenario, v_base
from optima_variables import Choice, database, normalizedWeight


# We consider the specification from the files `optima_variables.py` and `optima_specification.py`


# The objective of this laboratory is to use an estimated choice model to calculate various indicators. We start with
# the market shares and the revenues.


# We first create a data structure to store the value of an indicator and the bounds of its confidence interval.
class IndicatorTuple(NamedTuple):
    """Tuple storing the value of an indicator, and the bounds on its confidence interval."""

    value: float
    lower: float
    upper: float


# We build the loglikelihood function for the base model
logprob = loglogit(v_base, None, Choice)

# We estimate the parameters of the model
biogeme = BIOGEME(database, logprob)
model_name = 'optima_base'
biogeme.model_name = model_name
results: EstimationResults = biogeme.estimate(run_bootstrap=True, recycle=True)

# General statistics
print(results.short_summary())

# Estimated parameters
display(get_pandas_estimated_parameters(estimation_results=results))


# # Scenarios

# We are interested in the market shares of public transportation, as well as the revenues generated by this
# alternative.

# First, we calculate the market shares of each alternative for the current scenario, as well as the 90% confidence
# intervals.

# We write the expression for the choice probability of each alternative.
prob_PT = logit(v_base, None, 0)
prob_CAR = logit(v_base, None, 1)
prob_SM = logit(v_base, None, 2)


# We build a dictionary will all quantities to be simulated, that is, the choice probability of each alternative, as
# well as the weight of each observation.
simulate = {
    'weight': normalizedWeight,
    'Prob. PT': prob_PT,
    'Prob. car': prob_CAR,
    'Prob. SM': prob_SM,
}


# We perform the simulation and collect the simulated values
biosim = BIOGEME(database, simulate)
simulated_values = biosim.simulate(results.get_beta_values())


# We also calculate confidence intervals for the calculated quantities
b = results.get_betas_for_sensitivity_analysis()
left, right = biosim.confidence_intervals(b, 0.9)


# Market shares are calculated using the weighted mean of the individual probabilities. We first calculate the
# weighted quantities in an additional column of the data frame. Then, we calculate the mean of those columns.

# Alternative car
simulated_values['Weighted choice_prob. car'] = (
    simulated_values['weight'] * simulated_values['Prob. car']
)
left['Weighted choice_prob. car'] = left['weight'] * left['Prob. car']
right['Weighted choice_prob. car'] = right['weight'] * right['Prob. car']

marketShare_car = IndicatorTuple(
    value=simulated_values['Weighted choice_prob. car'].mean(),
    lower=left['Weighted choice_prob. car'].mean(),
    upper=right['Weighted choice_prob. car'].mean(),
)
print(
    f'Market share for car: {100*marketShare_car.value:.1f}% '
    f'[{100*marketShare_car.lower:.1f}%, '
    f'{100*marketShare_car.upper:.1f}%]'
)


# Alternative public transportation
simulated_values['Weighted choice_prob. PT'] = (
    simulated_values['weight'] * simulated_values['Prob. PT']
)
left['Weighted choice_prob. PT'] = left['weight'] * left['Prob. PT']
right['Weighted choice_prob. PT'] = right['weight'] * right['Prob. PT']

marketShare_PT = IndicatorTuple(
    value=simulated_values['Weighted choice_prob. PT'].mean(),
    lower=left['Weighted choice_prob. PT'].mean(),
    upper=right['Weighted choice_prob. PT'].mean(),
)
print(
    f'Market share for public transportation: {100*marketShare_PT.value:.1f}% '
    f'[{100*marketShare_PT.lower:.1f}%, '
    f'{100*marketShare_PT.upper:.1f}%]'
)


# Alternative slow modes
simulated_values['Weighted choice_prob. SM'] = (
    simulated_values['weight'] * simulated_values['Prob. SM']
)
left['Weighted choice_prob. SM'] = left['weight'] * left['Prob. SM']
right['Weighted choice_prob. SM'] = right['weight'] * right['Prob. SM']

marketShare_SM = IndicatorTuple(
    value=simulated_values['Weighted choice_prob. SM'].mean(),
    lower=left['Weighted choice_prob. SM'].mean(),
    upper=right['Weighted choice_prob. SM'].mean(),
)
print(
    f'Market share for slow modes: {100*marketShare_SM.value:.1f}% '
    f'[{100*marketShare_SM.lower:.1f}%, '
    f'{100*marketShare_SM.upper:.1f}%]'
)


# Now, consider various scenarios where the price of public transportation is multiplied for each individual by a
# scale factor. We advise you to build a Python function that takes the specification of the utility as an argument,
# and returns the market shares as well as the bounds of the confidence intervals.

# The function `scenario`, imported from `optima_specification`, takes a scale factor as an argument and returns a
# `ScenarioTuple` that contains
# - the expression for the utility functions,
# - the expression for the chosen alternative,
# - the expression for the cost:
#
# ```
# class ScenarioTuple(NamedTuple):
#     """Structure storing the expressions for a scenario"""
#
#     utilities: dict[int, Expression]
#     choice: Expression
#     cost: Expression
# ```

# Utility functions
print(scenario(factor=1.2).utilities)

# Choice
print(scenario(factor=1.2).choice)

# Cost
print(scenario(factor=1.2).cost)


# We consider scales ranging from 0 to 5, with a step of 0.2.
scales = np.arange(0, 5.2, 0.2)


# - Calculate the market share of public transportation, as well as the 90% confidence interval, for each scenario. Plot the results.
# - Calculate the market share of car, as well as the 90% confidence interval, for each scenario. Plot the results.
# - Calculate the revenues generated by public transportation, as well as the 90% confidence interval, for each scenario. Plot the results.