import larch  # !conda install larch #for estimation
import pandas as pd
import numpy as np
import yaml 
import larch.util.excel
import os

"""
Explanation: Estimating Joint and Non-Mandatory Tour Destination Choice
This notebook illustrates how to re-estimate both the joint tour destination and 
non-mandatory tour destination choice models.  These models share parameters and need
to be re-estimated together. This process 
includes running ActivitySim in estimation mode to read household travel survey files and write out
the estimation data bundles used in this notebook.  To review how to do so, please visit the other
notebooks in this directory.
Load libraries
End of explanation
"""


os.chdir('test')

"""
Explanation: We'll work in our test directory, where ActivitySim has saved the estimation data bundles.
End of explanation
"""


modelnames = ("non_mandatory_tour_destination", "joint_tour_destination")

from activitysim.estimation.larch import component_model
model, data = component_model(modelnames, return_data=True)

"""
Explanation: Load data and prep model for estimation
End of explanation
"""


type(model)

"""
Explanation: The resulting model is actually a larch ModelGroup, which exposes a similar 
API as a single model, and allows for joint estimation of all parameters across
all models in the group.
End of explanation
"""


data[0].coefficients

"""
Explanation: Review data loaded from EDB
Next we can review what was read the EDB, including the coefficients, model settings, utilities specification, and chooser and alternative data.
coefficients
End of explanation
"""


data[0].alt_values

data[1].alt_values

"""
Explanation: alt_values
End of explanation
"""


data[0].chooser_data

data[1].chooser_data

"""
Explanation: chooser_data
End of explanation
"""


data[0].landuse

"""
Explanation: landuse
End of explanation
"""


data[0].spec

"""
Explanation: spec
End of explanation
"""


data[0].size_spec

"""
Explanation: size_spec
End of explanation
"""


model.estimate(method='SLSQP')

"""
Explanation: Estimate
With the model setup for estimation, the next step is to estimate the model coefficients.  Make sure to use a sufficiently large enough household sample and set of zones to avoid an over-specified model, which does not have a numerically stable likelihood maximizing solution.  Larch has a built-in estimation methods including BHHH, and also offers access to more advanced general purpose non-linear optimizers in the scipy package, including SLSQP, which allows for bounds and constraints on parameters.  BHHH is the default and typically runs faster, but does not follow constraints on parameters.
End of explanation
"""


model.parameter_summary()

"""
Explanation: Estimated coefficients
End of explanation
"""


from activitysim.estimation.larch import update_coefficients, update_size_spec
result_dir = data[0].edb_directory/"estimated"

"""
Explanation: Output Estimation Results
End of explanation
"""


modelname = modelnames[0]

update_coefficients(
    model, data[0], result_dir,
    output_file=f"{modelname}_coefficients_revised.csv",
);

"""
Explanation: Write updated utility coefficients
The revised coefficients are written out as one file.  For the MTC
example model, these coefficients are written into the non-mandatory
tour model, and are re-used by the joint tour model.
End of explanation
"""


update_size_spec(
    model, data[0], result_dir, 
    output_file=f"{modelname}_size_terms.csv",
)

"""
Explanation: Write updated size coefficients
End of explanation
"""


model[0].to_xlsx(
    result_dir/f"{modelnames[0]}_model_estimation.xlsx", 
    data_statistics=False,
)
model[1].to_xlsx(
    result_dir/f"{modelnames[1]}_model_estimation.xlsx", 
    data_statistics=False,
);

"""
Explanation: Write the model estimation report, including coefficient t-statistic and log likelihood
Each model in the ModelGroup can have its own seperate estimation report, as the
common ModelGroup class does not yet implement a common report output.
End of explanation
"""


pd.read_csv(result_dir/f"{modelname}_coefficients_revised.csv")

pd.read_csv(result_dir/f"{modelname}_size_terms.csv")

"""
Explanation: Next Steps
The final step is to either manually or automatically copy the *_coefficients_revised.csv file and *_size_terms.csv file to the configs folder, rename them to *_coefficients.csv and destination_choice_size_terms.csv, and run ActivitySim in simulation mode.  Note that all the location
and desintation choice models share the same destination_choice_size_terms.csv input file, so if you
are updating all these models, you'll need to ensure that updated sections of this file for each model
are joined together correctly.
End of explanation
"""


