# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/correlations.ks.ipynb (unless otherwise specified).

__all__ = ['log', 'do_ks_dict']

# Cell

import logging

from scipy.stats import ks_2samp

from .utils import save_results
from ..data.utils import get_chart_map

log = logging.getLogger(__name__)

# Cell


def do_ks_dict(dict_baseline, dict_highlight, chart_map=None, model_level='dim', n_lags=0) -> dict:

    if not chart_map:
        chart_map = get_chart_map(dict_baseline, dict_highlight)

    # init some counters
    fit_success, fit_default, fit_fail, fit_manual = 0, 0, 0, 0

    # dict to collect results into
    results = {}

    for chart in chart_map:

        for dimension in chart_map[chart]:

            if chart not in dict_baseline:

                score = 1
                fit_manual += 1
                results = save_results(results, chart, dimension, score)

            elif chart not in dict_highlight:

                score = 1
                fit_manual += 1
                results = save_results(results, chart, dimension, score)

            elif dimension not in dict_baseline[chart]:

                score = 1
                fit_manual += 1
                results = save_results(results, chart, dimension, score)

            elif dimension not in dict_highlight[chart]:

                score = 1
                fit_manual += 1
                results = save_results(results, chart, dimension, score)

            else:

                data_baseline = dict_baseline[chart][dimension]
                data_highlight = dict_highlight[chart][dimension]

                try:
                    _, score = ks_2samp(data_baseline, data_highlight, mode='auto')
                    fit_success += 1
                    results = save_results(results, chart, dimension, score)
                except:
                    fit_fail += 1
                    continue

    # add summary to results
    results['summary'] = dict(
        fit_success=fit_success, fit_fail=fit_fail, fit_default=fit_default,
        model_level=model_level, n_lags=n_lags, fit_manual=fit_manual
    )

    return results

