import os
import larch  # !conda install larch -c conda-forge # for estimation
import pandas as pd

"""
Explanation: Estimating At-Work Subtour Frequency
This notebook illustrates how to re-estimate a single model component for ActivitySim.  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
"""


modelname = "atwork_subtour_frequency"

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

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


data.coefficients

"""
Explanation: Review data loaded from the EDB
The next step is to read the EDB, including the coefficients, model settings, utilities specification, and chooser and alternative data.
Coefficients
End of explanation
"""


data.spec

"""
Explanation: Utility specification
End of explanation
"""


data.chooser_data

"""
Explanation: Chooser data
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
result_dir = data.edb_directory/"estimated"
update_coefficients(
    model, data, result_dir,
    output_file=f"{modelname}_coefficients_revised.csv",
);

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


model.to_xlsx(
    result_dir/f"{modelname}_model_estimation.xlsx", 
    data_statistics=False,
)

"""
Explanation: Write the model estimation report, including coefficient t-statistic and log likelihood
End of explanation
"""


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

"""
Explanation: Next Steps
The final step is to either manually or automatically copy the *_coefficients_revised.csv file to the configs folder, rename it to *_coefficients.csv, and run ActivitySim in simulation mode.
End of explanation
"""


