"""File optima_specification.py

Michel Bierlaire, EPFL
Sun Aug 03 2025, 17:54:40

 Specification of a logit model, that will be estimated, and
 used for simulation.

"""

from typing import NamedTuple

from biogeme.expressions import Beta, Expression

from optima_variables import (
    Choice,
    CostCarCHF_scaled,
    MarginalCostPT,
    TimeCar_scaled,
    TimePT_scaled,
    distance_km_scaled,
    female,
    fulltime,
    male,
    notfulltime,
    unreportedGender,
)


class ScenarioTuple(NamedTuple):
    """Structure storing the expressions for a scenario"""

    utilities: dict[int, Expression]
    choice: Expression
    cost: Expression


# List of parameters to be estimated
asc_car = Beta('asc_car', 0, None, None, 0)
asc_pt = Beta('asc_pt', 0, None, None, 1)
asc_sm = Beta('asc_sm', 0, None, None, 0)
beta_time_fulltime = Beta('beta_time_fulltime', 0, None, None, 0)
beta_time_other = Beta('beta_time_other', 0, None, None, 0)
beta_dist_male = Beta('beta_dist_male', 0, None, None, 0)
beta_dist_female = Beta('beta_dist_female', 0, None, None, 0)
beta_dist_unreported = Beta('beta_dist_unreported', 0, None, None, 0)
beta_cost = Beta('beta_cost', 0, None, None, 0)


def scenario(factor: float = 1.0) -> ScenarioTuple:
    """Provide the model specification for a scenario with the price of
        public transportation is multiplied by a factor

    :param factor: factor that multiples the price of public transportation.

    :return: a tuple with the expressions of the utility function, the choice,
        and the cost.
    """
    MarginalCostScenario = MarginalCostPT * factor
    MarginalCostPT_scaled = MarginalCostScenario / 10
    # Definition of utility functions:
    v_pt = (
        asc_pt
        + beta_time_fulltime * TimePT_scaled * fulltime
        + beta_time_other * TimePT_scaled * notfulltime
        + beta_cost * MarginalCostPT_scaled
    )
    v_car = (
        asc_car
        + beta_time_fulltime * TimeCar_scaled * fulltime
        + beta_time_other * TimeCar_scaled * notfulltime
        + beta_cost * CostCarCHF_scaled
    )
    v_sm = (
        asc_sm
        + beta_dist_male * distance_km_scaled * male
        + beta_dist_female * distance_km_scaled * female
        + beta_dist_unreported * distance_km_scaled * unreportedGender
    )

    # Associate utility functions with the numbering of alternatives
    v = {0: v_pt, 1: v_car, 2: v_sm}

    return ScenarioTuple(utilities=v, choice=Choice, cost=MarginalCostScenario)


# The base scenario is the default
base_scenario = scenario()
v_base = base_scenario.utilities