# TEST
import larch.numba as lx
from pytest import approx

import larch.numba as lx
from larch import P, X

"""
Explanation: 203: Exampville Destination Choice
Welcome to Exampville, the best simulated town in this here part of the internet!
Exampville is a demonstration provided with Larch that walks through some of the 
data and tools that a transportation planner might use when building a travel model.
End of explanation
"""


hh, pp, tour, skims, emp = lx.example(200, ['hh', 'pp', 'tour', 'skims', 'emp'])

"""
Explanation: In this example notebook, we will walk through the estimation of a tour 
destination choice model.  First, let's load the data files from
our example.
End of explanation
"""


logsums_file = lx.example(202, output_file='logsums.zarr.zip')
logsums = lx.DataArray.from_zarr('logsums.zarr.zip')

"""
Explanation: For this destination choice model, we'll want to use the mode choice
logsums we calculated previously from the mode choice estimation,
but we'll use these values as fixed input data instead of a modeled value.
We can load these logsums from the file in which they were saved. 
For this example, we can indentify that file using the larch.example 
function, which will automatically rebuild the file if it doesn't exists.
In typical applications, a user would generally just give the filename 
as a string and ensure manually that the file exists before loading it.
End of explanation
"""


ca = lx.Dataset.construct(
    {'logsum': logsums},
    caseid='TOURID',
    alts=skims.TAZ_ID,
)
ca

"""
Explanation: Preprocessing
The alternatives in
the destinations model are much more regular than in the mode choice 
model, as every observation will have a similar set of alternatives
and the utility function for each of those alternatives will share a 
common functional form.  We'll leverage this by using idca format 
arrays in our DataTree to make data management simpler.  
The base array we'll start with will have two dimensions, cases and
alternatives, anc be formed from the logsums we loaded above.
End of explanation
"""


emp.info()

"""
Explanation: For our destination choice model, we'll also want to use employment data.
This data, as included in our example, has unique 
values only by alternative and not by caseid, so there are only
40 unique rows.
(This kind of structure is common for destination choice models.)
End of explanation
"""


tree = lx.DataTree(
    base=ca,
    tour=tour.rename_axis(index='TOUR_ID'),
    hh=hh.set_index("HHID"),
    person=pp.set_index('PERSONID'),
    emp=emp,
    skims=lx.Dataset.construct.from_omx(skims),
    relationships=(
        "base.TAZ_ID @ emp.TAZ",
        "base.TOURID -> tour.TOUR_ID",
        "tour.HHID @ hh.HHID",
        "tour.PERSONID @ person.PERSONID",
        "hh.HOMETAZ @ skims.otaz",
        "base.TAZ_ID @ skims.dtaz"
    ),
)

"""
Explanation: Then we bundle all our raw data into a DataTree structure, 
which is used to collect the right data for estimation.  The
Larch DataTree is a slightly augmented version of the regular
sharrow.DataTree.
End of explanation
"""


m = lx.Model(datatree=tree)
m.title = "Exampville Work Tour Destination Choice v1"

m.quantity_ca = (
        + P.EmpRetail_HighInc * X('RETAIL_EMP * (INCOME>50000)')
        + P.EmpNonRetail_HighInc * X('NONRETAIL_EMP') * X("INCOME>50000")
        + P.EmpRetail_LowInc * X('RETAIL_EMP') * X("INCOME<=50000")
        + P.EmpNonRetail_LowInc * X('NONRETAIL_EMP') * X("INCOME<=50000")
)

m.quantity_scale = P.Theta


m.utility_ca = (
    + P.logsum * X.logsum
    + P.distance * X.AUTO_DIST
)

m.choice_co_code = "tour.DTAZ"

m.lock_values(
    EmpRetail_HighInc=0,
    EmpRetail_LowInc=0,
)

"""
Explanation: Model Definition
Now we can define our choice model, using data from the tree as appropriate.
End of explanation
"""


m.loglike()

# TEST
assert m.loglike() == approx(-28238.336880999712)

m.maximize_loglike()

# TEST
result = _
assert result.loglike == approx(-25157.72676137825)
assert result.success
assert result.method == 'slsqp'
assert result.n_cases == 7564
assert result.iteration_number == 13
assert result.logloss == approx(3.325981856342973)
import pandas as pd
pd.testing.assert_series_equal(
    result.x.sort_index(),
    pd.Series({
        'EmpNonRetail_HighInc': 1.3639677867685673,
        'EmpNonRetail_LowInc': -0.8813930080108792,
        'EmpRetail_HighInc': 0.0,
        'EmpRetail_LowInc': 0.0,
        'Theta': 0.7493698789904475,
        'distance': -0.04182074282571401,
        'logsum': 1.0208113241683572,
    }).sort_index(),
    rtol=1e-3,
)

m.calculate_parameter_covariance()

"""
Explanation: Model Estimation
End of explanation
"""


m.distribution_on_idca_variable('AUTO_DIST')

"""
Explanation: Model Visualization
For destination choice and similar type models, it might be beneficial to
review the observed and modeled choices, and the relative distribution of
these choices across different factors.  For example, we would probably want
to see the distribution of travel distance.  The Model object includes
a built-in method to create this kind of visualization.
End of explanation
"""


m.distribution_on_idca_variable('AUTO_DIST', bins=40, range=(0,10))

"""
Explanation: The distribution_on_idca_variable has a variety of options,
for example to control the number and range of the histogram bins:
End of explanation
"""


m.distribution_on_idca_variable(
    'AUTO_DIST',
    style='kde',
)

"""
Explanation: Alternatively, the histogram style can be swapped out for a smoothed kernel density
function:
End of explanation
"""


m.distribution_on_idca_variable(
    'AUTO_DIST',
    xlabel="Distance (miles)",
    bins=26,
    subselector='INCOME<10000',
    range=(0,13),
    header='Destination Distance, Very Low Income (<$10k) Households',
)

"""
Explanation: Subsets of the observations can be pulled out, to observe the 
distribution conditional on other idco factors, like income.
End of explanation
"""


report = lx.Reporter(title=m.title)

report << '# Parameter Summary' << m.parameter_summary()

# TEST
assert m.parameter_summary().data.to_markdown() == '''
|                      |   Value | Std Err   | t Stat   | Signif   |   Null Value | Constrained   |
|:---------------------|--------:|:----------|:---------|:---------|-------------:|:--------------|
| EmpNonRetail_HighInc |  1.36   | 0.256     | 5.32     | ***      |            0 |               |
| EmpNonRetail_LowInc  | -0.881  | 0.0791    | -11.14   | ***      |            0 |               |
| EmpRetail_HighInc    |  0      | NA        | NA       |          |            0 | fixed value   |
| EmpRetail_LowInc     |  0      | NA        | NA       |          |            0 | fixed value   |
| Theta                |  0.749  | 0.0152    | -16.45   | ***      |            1 |               |
| distance             | -0.0418 | 0.0107    | -3.90    | ***      |            0 |               |
| logsum               |  1.02   | 0.0317    | 32.16    | ***      |            0 |               |
'''[1:-1]

report << "# Estimation Statistics" << m.estimation_statistics()

report << "# Utility Functions" << m.utility_functions()

"""
Explanation: Save and Report Model
End of explanation
"""


figure = m.distribution_on_idca_variable(
    'AUTO_DIST', 
    xlabel="Distance (miles)",
    style='kde',
    header='Destination Distance',
)
report << "# Visualization"
report << figure

report.save(
    'exampville_dest_choice.html',
    overwrite=True,
    metadata=m,
)

"""
Explanation: The figures shown above can also be inserted directly into reports.
End of explanation
"""


