

"""
Simulating/estimating discrete choice experiments within sfctools
NOTE before starting this script, download the swissmetro data from
https://transp-or.epfl.ch/pythonbiogeme/examples/swissmetro/swissmetro.dat
"""

import pandas as pd
import sys
import numpy as np


if __name__ == "__main__":

    from sfctools.ml_models.nested_logit import MultinomialNestedLogitModel as MNL

    # ------------------------------
    # PREDICT CHOICES USING THEORETICAL MODEL

    # limit data length
    nesting = {
        "1": ["feature1"],
        "2": ["feature2"],
    }

    model_pred = MNL(nesting, numeric_indexing=False)

    # params_true = np.array([0.0, .3, -0.1])
    params_true = np.array([0.07, 0.08, -.1])
    # model.visualize_tree()  # < visualizes the model tree as networkx graph

    # print(model_pred.param_vector)
    # print([id(i) for i in model_pred.param_vector.to_serial()])
    # # sys.exit()
    # for k, v in model_pred.probabilities.items():
    #     print(k, "p=", v)
    # sys.exit()

    # predict values
    model_pred.param_vector.set_values(params_true)

    # invent some random data
    N_obs = 50_000
    data = pd.DataFrame({
        "feature1": [np.round(xi, 1) for xi in np.random.uniform(0.01, 10.1, N_obs)],
        "feature2": [np.round(xi, 1) for xi in np.random.uniform(0.01, 10.1, N_obs)]
    })

    # assign choice outputs from the model itself to the artificial data
    choices = model_pred.predict(data, noise=0.0)  # add a little bit of noise
    data["choice"] = choices

    # print("choice probabilities")
    # for c in data["choice"].unique():
    #     prob = len(data[data["choice"] == c])/len(data)
    #     print(f"{c}: p={prob:.4f}")
    # print("\n")

    # ------------------------------
    # TRY TO FIT MODEL

    # set parameter bounds
    bounds = {
        # normalize beat_1_feature1 to 0 as is alternative-specific feature
        'beta_1_feature1': (-10, 10),
        'beta_2_feature2': (-200, 200),
        'asc_2': (-200, 200)
    }

    # setup a model with same nesting
    model_fit = MNL(nesting, numeric_indexing=False)
    # , equal_params={
    #    "beta_2_feature2": "beta_1_feature1"})

    # print(model_fit.param_vector)
    # for k, v in model_fit.probabilities.items():
    #     print(k, "P=", v)

    params_guess = np.array([-0.5, -0.5, 0.1])
    model_fit.param_vector.set_values(params_guess)

    # # fit the model
    result, stats = model_fit.fit(
        data, sampling_method='lhs', num_samples_per_param=1000, scale_attrs=False,
        bounds=bounds, optimize_globally=True,  # opt_method="SLSQP",
        regularization="l2", alpha_2=10.0)

    model_fit.param_vector.set_values(result.x)
    # add a little bit of noise
    choices = model_fit.predict(data, noise=0.0)
    data["choice_fit"] = choices

    # print("\n")
    # print("choice probabilities")
    # for c in data["choice_fit"].unique():
    #     prob = len(data[data["choice_fit"] == c])/len(data)
    #     print(f"{c}: p={prob:.4f}")
    # print("\n")

    del data["choice_fit"]

    # tweak asc
    params_tweak, r2 = model_fit.tweak_asc(data)
    # print("best params after tweak", params_tweak)
