"""File spec_optima.py

Michel Bierlaire, EPFL
Mon Aug 04 2025, 10:25:52

"""

from biogeme.expressions import Beta, Variable
from biogeme.models import boxcox
from biogeme.segmentation import DiscreteSegmentationTuple, Segmentation

# Definition of the variables in the data file
Id = Variable('Id')
MarginalCostPT = Variable('MarginalCostPT')
WaitingTimePT = Variable('WaitingTimePT')
CostCarCHF = Variable('CostCarCHF')
NbTransf = Variable('NbTransf')
distance_km = Variable('distance_km')
TimePT = Variable('TimePT')
TimeCar = Variable('TimeCar')
OccupStat = Variable('OccupStat')
LangCode = Variable('LangCode')
CarAvail = Variable('CarAvail')
Education = Variable('Education')
TripPurpose = Variable('TripPurpose')
Prob0 = Variable('Prob0')
Prob1 = Variable('Prob1')
Prob2 = Variable('Prob2')
Choice = Variable('Choice')

TimePT_scaled = TimePT / 200
TimeCar_scaled = TimeCar / 200
MarginalCostPT_scaled = MarginalCostPT / 10
CostCarCHF_scaled = CostCarCHF / 10
distance_km_scaled = distance_km / 5

car_avail_segmentation = DiscreteSegmentationTuple(
    variable=CarAvail, mapping={1: 'car_avail', 3: 'car_unavail'}
)

language_segmentation = DiscreteSegmentationTuple(
    variable=LangCode, mapping={1: 'french', 2: 'german'}
)

occup_segmentation = DiscreteSegmentationTuple(
    variable=OccupStat, mapping={1: 'full_time', 2: 'part_time', 3: 'others'}
)

purpose_segmentation = DiscreteSegmentationTuple(
    variable=TripPurpose, mapping={1: 'work', 2: 'not_work'}
)

education_segmentation = DiscreteSegmentationTuple(
    variable=Education,
    mapping={
        3: 'vocational',
        4: 'high_school',
        6: 'higher_education',
        7: 'university',
    },
)

asc_pt = Beta('asc_pt', 0, None, None, 0)
segmented_asc_pt = Segmentation(
    asc_pt,
    [car_avail_segmentation, language_segmentation],
).segmented_beta()

asc_car = Beta('asc_car', 0, None, None, 0)
segmented_asc_car = Segmentation(
    asc_car,
    [language_segmentation],
).segmented_beta()

beta_time = Beta('beta_time', 0, None, None, 0)
segmented_beta_time = Segmentation(beta_time, [occup_segmentation]).segmented_beta()

beta_cost_pt = Beta('beta_cost_pt', 0, None, None, 0)
beta_cost_car = Beta('beta_cost_car', 0, None, None, 0)

beta_waiting = Beta('beta_waiting', 0, None, None, 0)
segmented_beta_waiting = Segmentation(
    beta_waiting, [purpose_segmentation]
).segmented_beta()

beta_dist = Beta('beta_dist', 0, None, None, 0)
segmented_beta_dist = Segmentation(beta_dist, [education_segmentation]).segmented_beta()

LAMBDA_COST = 0.3214999879822265

v_pt = (
    segmented_asc_pt
    + segmented_beta_time * TimePT_scaled
    + beta_cost_pt * boxcox(MarginalCostPT_scaled, LAMBDA_COST)
    + segmented_beta_waiting * WaitingTimePT**0.5
)

v_car = (
    segmented_asc_car
    + segmented_beta_time * TimeCar_scaled
    + beta_cost_car * boxcox(CostCarCHF_scaled, LAMBDA_COST)
)

v_sm = segmented_beta_dist * distance_km_scaled


v = {0: v_pt, 1: v_car, 2: v_sm}

av = {0: 1, 1: CarAvail != 3, 2: 1}