#!/usr/bin/env python
# Copyright (c) 2017 The WebRTC project authors. All Rights Reserved.
#
# Use of this source code is governed by a BSD-style license
# that can be found in the LICENSE file in the root of the source
# tree. An additional intellectual property rights grant can be found
# in the file PATENTS.  All contributing project authors may
# be found in the AUTHORS file in the root of the source tree.
"""Finds the APM configuration that maximizes a provided metric by
parsing the output generated apm_quality_assessment.py.
"""

from __future__ import division

import collections
import logging
import os

import quality_assessment.data_access as data_access
import quality_assessment.collect_data as collect_data


def _InstanceArgumentsParser():
    """Arguments parser factory. Extends the arguments from 'collect_data'
  with a few extra for selecting what parameters to optimize for.
  """
    parser = collect_data.InstanceArgumentsParser()
    parser.description = (
        'Rudimentary optimization of a function over different parameter'
        'combinations.')

    parser.add_argument(
        '-n',
        '--config_dir',
        required=False,
        help=('path to the folder with the configuration files'),
        default='apm_configs')

    parser.add_argument('-p',
                        '--params',
                        required=True,
                        nargs='+',
                        help=('parameters to parse from the config files in'
                              'config_dir'))

    parser.add_argument(
        '-z',
        '--params_not_to_optimize',
        required=False,
        nargs='+',
        default=[],
        help=('parameters from `params` not to be optimized for'))

    return parser


def _ConfigurationAndScores(data_frame, params, params_not_to_optimize,
                            config_dir):
    """Returns a list of all configurations and scores.

  Args:
    data_frame: A pandas data frame with the scores and config name
                returned by _FindScores.
    params: The parameter names to parse from configs the config
            directory

    params_not_to_optimize: The parameter names which shouldn't affect
                            the optimal parameter
                            selection. E.g., fixed settings and not
                            tunable parameters.

    config_dir: Path to folder with config files.

  Returns:
    Dictionary of the form
    {param_combination: [{params: {param1: value1, ...},
                          scores: {score1: value1, ...}}]}.

    The key `param_combination` runs over all parameter combinations
    of the parameters in `params` and not in
    `params_not_to_optimize`. A corresponding value is a list of all
    param combinations for params in `params_not_to_optimize` and
    their scores.
  """
    results = collections.defaultdict(list)
    config_names = data_frame['apm_config'].drop_duplicates().values.tolist()
    score_names = data_frame['eval_score_name'].drop_duplicates(
    ).values.tolist()

    # Normalize the scores
    normalization_constants = {}
    for score_name in score_names:
        scores = data_frame[data_frame.eval_score_name == score_name].score
        normalization_constants[score_name] = max(scores)

    params_to_optimize = [p for p in params if p not in params_not_to_optimize]
    param_combination = collections.namedtuple("ParamCombination",
                                               params_to_optimize)

    for config_name in config_names:
        config_json = data_access.AudioProcConfigFile.Load(
            os.path.join(config_dir, config_name + ".json"))
        scores = {}
        data_cell = data_frame[data_frame.apm_config == config_name]
        for score_name in score_names:
            data_cell_scores = data_cell[data_cell.eval_score_name ==
                                         score_name].score
            scores[score_name] = sum(data_cell_scores) / len(data_cell_scores)
            scores[score_name] /= normalization_constants[score_name]

        result = {'scores': scores, 'params': {}}
        config_optimize_params = {}
        for param in params:
            if param in params_to_optimize:
                config_optimize_params[param] = config_json['-' + param]
            else:
                result['params'][param] = config_json['-' + param]

        current_param_combination = param_combination(**config_optimize_params)
        results[current_param_combination].append(result)
    return results


def _FindOptimalParameter(configs_and_scores, score_weighting):
    """Finds the config producing the maximal score.

  Args:
    configs_and_scores: structure of the form returned by
                        _ConfigurationAndScores

    score_weighting: a function to weight together all score values of
                     the form [{params: {param1: value1, ...}, scores:
                                {score1: value1, ...}}] into a numeric
                     value
  Returns:
    the config that has the largest values of `score_weighting` applied
    to its scores.
  """

    min_score = float('+inf')
    best_params = None
    for config in configs_and_scores:
        scores_and_params = configs_and_scores[config]
        current_score = score_weighting(scores_and_params)
        if current_score < min_score:
            min_score = current_score
            best_params = config
            logging.debug("Score: %f", current_score)
            logging.debug("Config: %s", str(config))
    return best_params


def _ExampleWeighting(scores_and_configs):
    """Example argument to `_FindOptimalParameter`
  Args:
    scores_and_configs: a list of configs and scores, in the form
                        described in _FindOptimalParameter
  Returns:
    numeric value, the sum of all scores
  """
    res = 0
    for score_config in scores_and_configs:
        res += sum(score_config['scores'].values())
    return res


def main():
    # Init.
    # TODO(alessiob): INFO once debugged.
    logging.basicConfig(level=logging.DEBUG)
    parser = _InstanceArgumentsParser()
    args = parser.parse_args()

    # Get the scores.
    src_path = collect_data.ConstructSrcPath(args)
    logging.debug('Src path <%s>', src_path)
    scores_data_frame = collect_data.FindScores(src_path, args)
    all_scores = _ConfigurationAndScores(scores_data_frame, args.params,
                                         args.params_not_to_optimize,
                                         args.config_dir)

    opt_param = _FindOptimalParameter(all_scores, _ExampleWeighting)

    logging.info('Optimal parameter combination: <%s>', opt_param)
    logging.info('It\'s score values: <%s>', all_scores[opt_param])


if __name__ == "__main__":
    main()