import larch
import pandas as pd
from larch.roles import P,X

"""
Explanation: 107: Latent Class Models
In this example, we will replicate the latent class example model
from Biogeme.
End of explanation
"""


from larch import data_warehouse
raw = pd.read_csv(larch.data_warehouse.example_file('swissmetro.csv.gz'))

"""
Explanation: The swissmetro dataset used in this example is conveniently bundled with Larch,
accessible using the data_warehouse module.  We'll load this file using 
the pandas read_csv command.
End of explanation
"""


raw.head()

"""
Explanation: We can inspect a few rows of data to see what we have using the head method.
End of explanation
"""


raw['SM_COST'] = raw['SM_CO'] * (raw["GA"]==0) 

"""
Explanation: The Biogeme code includes a variety of commands to manipulate the data
and create new variables.  Because Larch sits on top of pandas, a reasonable
method to create new variables is to just create new columns in the 
source pandas.DataFrame in the usual manner for any DataFrame.
End of explanation
"""


raw['TRAIN_COST'] = raw.eval("TRAIN_CO * (GA == 0)") 

"""
Explanation: You can also use the eval method of pandas DataFrames.
This method takes an expression as a string
and evaluates it within a namespace that has already loaded the 
column names as variables.
End of explanation
"""


raw['TRAIN_COST_SCALED'] = raw['TRAIN_COST'] / 100
raw['TRAIN_TT_SCALED'] = raw['TRAIN_TT'] / 100

raw['SM_COST_SCALED'] = raw.eval('SM_COST / 100')
raw['SM_TT_SCALED'] = raw['SM_TT'] / 100

raw['CAR_CO_SCALED'] = raw['CAR_CO'] / 100
raw['CAR_TT_SCALED'] = raw['CAR_TT'] / 100

raw['CAR_AV_SP'] = raw.eval("CAR_AV * (SP!=0)")
raw['TRAIN_AV_SP'] = raw.eval("TRAIN_AV * (SP!=0)")

"""
Explanation: This can allow for writing data 
expressions more succinctly, as long as all your variable names 
are strings that can also be the names of variables in Python.
If this isn't the case (e.g., if any variable names have spaces
in the name) you'll be better off if you stay away from this 
feature.
We can mix and match between these two method to create new
columns in any DataFrame as needed.
End of explanation
"""


keep = raw.eval("PURPOSE in (1,3) and CHOICE != 0")

"""
Explanation: Removing some observations can also be done directly using pandas.
Here we identify a subset of observations that we want to keep.
End of explanation
"""


dfs = larch.DataFrames(raw[keep], alt_codes=[1,2,3])

"""
Explanation: You may note that we don't assign this value to a column within the
raw DataFrame.  This is perfectly acceptable, as the output from
the eval method is just a normal pandas.Series, like any other
single column output you might expect to get from a pandas method.
When you've created the data you need, you can pass the dataframe to
the larch.DataFrames constructor.  Since the swissmetro data is in
idco format, we'll need to explicitly identify the alternative 
codes as well.
End of explanation
"""


dfs.info()

"""
Explanation: The info method of the DataFrames object gives a short summary
of the contents.
End of explanation
"""


dfs.info(verbose=True)

"""
Explanation: A longer summary is available by setting verbose to True.
End of explanation
"""


dfs.data_co.info()

"""
Explanation: You may have noticed that the info summary notes that this data is "not computation-ready".
That's because some of the data columns are stored as integers, which can be observed by
inspecting the info on the data_co dataframe.
End of explanation
"""


m1 = larch.Model(dataservice=dfs)
m1.availability_co_vars = {
    1: "TRAIN_AV_SP",
    2: "SM_AV",
    3: "CAR_AV_SP",
}
m1.choice_co_code = 'CHOICE'

m1.utility_co[1] = P("ASC_TRAIN") + X("TRAIN_COST_SCALED") * P("B_COST")
m1.utility_co[2] = X("SM_COST_SCALED") * P("B_COST")
m1.utility_co[3] = P("ASC_CAR") + X("CAR_CO_SCALED") * P("B_COST")


m2 = larch.Model(dataservice=dfs)
m2.availability_co_vars = {
    1: "TRAIN_AV_SP",
    2: "SM_AV",
    3: "CAR_AV_SP",
}
m2.choice_co_code = 'CHOICE'

m2.utility_co[1] = P("ASC_TRAIN") + X("TRAIN_TT_SCALED") * P("B_TIME") + X("TRAIN_COST_SCALED") * P("B_COST")
m2.utility_co[2] = X("SM_TT_SCALED") * P("B_TIME") + X("SM_COST_SCALED") * P("B_COST")
m2.utility_co[3] = P("ASC_CAR") + X("CAR_TT_SCALED") * P("B_TIME") + X("CAR_CO_SCALED") * P("B_COST")

"""
Explanation: When computations are run, we'll need all the data to be in float format, but Larch knows this and will 
handle it for you later.
Class Model Setup
Having prepped our data, we're ready to set up discrete choices models 
for each class in the latent class model.  We'll reproduce the Biogeme
example exactly here, as a technology demonstation.  Each of two classes
will be set up with a simple MNL model.
End of explanation
"""


mk = larch.Model()
mk.utility_co[2] = P("W_OTHER")

"""
Explanation: Class Membership Model
For Larch, the class membership model will be set up as yet another discrete choice model.
In this case, the choices are not the ultimate choices, but instead are the latent classes.
To remain consistent with the Biogeme example, we'll set up this model with only a single
constant that determines class membership.  Unlike Biogeme, this class membership will
be represented with an MNL model, not a simple direct probability.
End of explanation
"""


from larch.model.latentclass import LatentClassModel
m = LatentClassModel(mk, {1:m1, 2:m2})

"""
Explanation: The utility function of the first class isn't written here, which means it will implicitly
be set as 0.
Latent Class Model
Now we're ready to create the latent class model itself, by assembling the components
we created above.  The constructor for the LatentClassModel takes two arguments,
a class membership model, and a dictionary of class models, where the keys in the 
dictionary correspond to the identifying codes from the utility functions we wrote
for the class membership model.
End of explanation
"""


m.load_data()

m.dataframes.info(verbose=1)

"""
Explanation: The we'll load the data needed for our models using the load_data method.
This step will assemble the data needed, and convert it to floating point
format as required.
End of explanation
"""


result = m.maximize_loglike()

result

"""
Explanation: Only the data actually needed by the models has been converted, which may help
keep memory usage down on larger models.  You may also note that the loaded
dataframes no longer reports that it is "not computational-ready".
To estimate the model, we'll use the maximize_loglike method.  When run
in Jupyter, a live-view report of the parmeters and log likelihood is displayed.
End of explanation
"""


m.loglike_null()

"""
Explanation: To complete our analysis, we can compute the log likelihood at "null" parameters.
End of explanation
"""


m.calculate_parameter_covariance()

m.covariance_matrix

m.robust_covariance_matrix

"""
Explanation: And the parameter covariance matrixes.
End of explanation
"""


report = larch.Reporter("Latent Class Example")

"""
Explanation: Reporting Results
And then generate a report of the estimation statistics.  Larch includes a Reporter class 
to help you assemble a report containing the relevant output you want.
End of explanation
"""


report << "# Parameter Estimates"

"""
Explanation: Pipe into the report section headers in markdown format (use one hash for top level
headings, two hashes for lower levels, etc.)
End of explanation
"""


report << m.pf

"""
Explanation: You can also pipe in dataframes directly, include the pf parameter frame from the model.
End of explanation
"""


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

report << "# Parameter Covariance"
report << "## Typical Parameter Covariance"
report << m.covariance_matrix
report << "## Robust Parameter Covariance"
report << m.robust_covariance_matrix

report << "# Utility Functions"
report << "## Class 1"
report << "### Formulae"
report << m1.utility_functions(resolve_parameters=False)
report << "### Final Estimated Values"
report << m1.utility_functions(resolve_parameters=True)

report << "## Class 2"
report << "### Formulae"
report << m1.utility_functions(resolve_parameters=False)
report << "### Final Estimated Values"
report << m1.utility_functions(resolve_parameters=True)

"""
Explanation: And a selection of pre-formatted summary sections.
End of explanation
"""


report.save('latent-class-example-report.html', overwrite=True)

"""
Explanation: In addition to reviewing report sections in a Jupyter notebook, the 
entire report can be saved to an HTML file.
End of explanation
"""


