# ActivitySim
# See full license in LICENSE.txt.
import logging

import numpy as np

from activitysim.core import tracing
from activitysim.core import config
from activitysim.core import pipeline
from activitysim.core import simulate
from activitysim.core import inject
from activitysim.core import expressions

from activitysim.abm.models.util import estimation

logger = logging.getLogger("activitysim")


@inject.step()
def work_from_home(
        persons_merged, persons,
        chunk_size, trace_hh_id):
    """
    This model predicts whether a person (worker) works from home. The output
    from this model is TRUE (if works from home) or FALSE (works away from home).
    """

    trace_label = 'work_from_home'
    model_settings_file_name = 'work_from_home.yaml'

    choosers = persons_merged.to_frame()
    model_settings = config.read_model_settings(model_settings_file_name)
    chooser_filter_column_name = model_settings.get('CHOOSER_FILTER_COLUMN_NAME')
    choosers = choosers[choosers[chooser_filter_column_name]]
    logger.info("Running %s with %d persons", trace_label, len(choosers))

    estimator = estimation.manager.begin_estimation('work_from_home')

    constants = config.get_model_constants(model_settings)
    work_from_home_alt = model_settings['WORK_FROM_HOME_ALT']

    # - preprocessor
    preprocessor_settings = model_settings.get('preprocessor', None)
    if preprocessor_settings:

        locals_d = {}
        if constants is not None:
            locals_d.update(constants)

        expressions.assign_columns(
            df=choosers,
            model_settings=preprocessor_settings,
            locals_dict=locals_d,
            trace_label=trace_label)

    model_spec = simulate.read_model_spec(file_name=model_settings['SPEC'])
    coefficients_df = simulate.read_model_coefficients(model_settings)

    nest_spec = config.get_logit_model_settings(model_settings)

    if estimator:
        estimator.write_model_settings(model_settings, model_settings_file_name)
        estimator.write_spec(model_settings)
        estimator.write_coefficients(coefficients_df)
        estimator.write_choosers(choosers)

    # - iterative single process what-if adjustment if specified
    iterations = model_settings.get('WORK_FROM_HOME_ITERATIONS', 1)
    iterations_chooser_filter = model_settings.get('WORK_FROM_HOME_CHOOSER_FILTER', None)
    iterations_coefficient_constant = model_settings.get('WORK_FROM_HOME_COEFFICIENT_CONSTANT', None)
    iterations_target_percent = model_settings.get('WORK_FROM_HOME_TARGET_PERCENT', None)
    iterations_target_percent_tolerance = model_settings.get('WORK_FROM_HOME_TARGET_PERCENT_TOLERANCE', None)

    for iteration in range(iterations):

        logger.info("Running %s with %d persons iteration %d", trace_label, len(choosers), iteration)

        # re-read spec to reset substitution
        model_spec = simulate.read_model_spec(file_name=model_settings['SPEC'])
        model_spec = simulate.eval_coefficients(model_spec, coefficients_df, estimator)

        choices = simulate.simple_simulate(
            choosers=choosers,
            spec=model_spec,
            nest_spec=nest_spec,
            locals_d=constants,
            chunk_size=chunk_size,
            trace_label=trace_label,
            trace_choice_name='work_from_home',
            estimator=estimator)

        if iterations_target_percent is not None:
            choices_for_filter = choices[choosers[iterations_chooser_filter]]

            current_percent = ((choices_for_filter == work_from_home_alt).sum() / len(choices_for_filter))
            logger.info("Running %s iteration %i choosers %i current percent %f target percent %f",
                        trace_label, iteration, len(choices_for_filter), current_percent, iterations_target_percent)

            if current_percent <= (iterations_target_percent +
                                   iterations_target_percent_tolerance
                                   ) and current_percent >= (iterations_target_percent -
                                                             iterations_target_percent_tolerance):
                logger.info("Running %s iteration %i converged with coefficient %f", trace_label, iteration,
                            coefficients_df.value[iterations_coefficient_constant])
                break

            else:
                new_value = np.log(iterations_target_percent /
                                   np.maximum(current_percent, 0.0001)
                                   ) + coefficients_df.value[iterations_coefficient_constant]
                coefficients_df.value[iterations_coefficient_constant] = new_value
                logger.info("Running %s iteration %i new coefficient for next iteration %f",
                            trace_label, iteration, new_value)
                iteration = iteration + 1

    choices = (choices == work_from_home_alt)

    if estimator:
        estimator.write_choices(choices)
        choices = estimator.get_survey_values(choices, 'persons', 'work_from_home')
        estimator.write_override_choices(choices)
        estimator.end_estimation()

    persons = persons.to_frame()
    persons['work_from_home'] = choices.reindex(persons.index).fillna(0).astype(bool)
    persons['is_out_of_home_worker'] = persons[chooser_filter_column_name] & ~persons['work_from_home']

    pipeline.replace_table("persons", persons)

    tracing.print_summary('work_from_home', persons.work_from_home, value_counts=True)

    if trace_hh_id:
        tracing.trace_df(persons,
                         label=trace_label,
                         warn_if_empty=True)
