# mmpdb - matched molecular pair database generation and analysis
#
# Copyright (c) 2015-2017, F. Hoffmann-La Roche Ltd.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#    * Redistributions of source code must retain the above copyright
#      notice, this list of conditions and the following disclaimer.
#    * Redistributions in binary form must reproduce the above
#      copyright notice, this list of conditions and the following
#      disclaimer in the documentation and/or other materials provided
#      with the distribution.
#    * Neither the name of F. Hoffmann-La Roche Ltd. nor the names of
#      its contributors may be used to endorse or promote products
#      derived from this software without specific prior written
#      permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from __future__ import print_function, absolute_import, division

import sys
import os

from . import command_support
from . import index_algorithm
from . import fragment_io
from . import fragment_types
from . import properties_io
from ._compat import open_universal

try:
    import psutil
except ImportError:
    psutil = None


def get_fragment_filter_from_args(parser, args):
    min_variable_heavies = args.min_variable_heavies
    max_variable_heavies = args.max_variable_heavies
    min_variable_ratio = args.min_variable_ratio
    max_variable_ratio = args.max_variable_ratio

    if max_variable_heavies == "none":  # Push this into the argument parser?
        max_variable_heavies = None

    if min_variable_ratio is not None and max_variable_ratio is not None:
        if min_variable_ratio > max_variable_ratio:
            parser.error("--min-variable-ratio must not be larger than --max-variable-ratio")

    if min_variable_heavies is not None and max_variable_heavies is not None:
        if min_variable_heavies > max_variable_heavies:
            parser.error("--min-variable-heavies must not be larger than --max-variable-heavies")

    filters = []
    if min_variable_heavies is not None:
        filters.append(index_algorithm.MinVariableHeaviesFilter(min_variable_heavies))
    if max_variable_heavies is not None:
        filters.append(index_algorithm.MaxVariableHeaviesFilter(max_variable_heavies))
    if min_variable_ratio is not None:
        filters.append(index_algorithm.MinVariableRatioFilter(min_variable_ratio))
    if max_variable_ratio is not None:
        filters.append(index_algorithm.MaxVariableRatioFilter(max_variable_ratio))

    if not filters:
        # It's easier to have 0 filters than to make a special do-nothing filter.
        return index_algorithm.MultipleFilters([])
    elif len(filters) == 1:
        return filters[0]
    else:
        return index_algorithm.MultipleFilters(filters)


if psutil is None:
    def get_memory_use():
        return 0
else:
    _process = psutil.Process(os.getpid())


    def get_memory_use():
        info = _process.memory_info()
        return info.rss  # or info.vms?


def human_memory(n):
    if n < 1024:
        return "%d B" % (n,)
    for unit, denom in (("KB", 1024), ("MB", 1024 ** 2),
                        ("GB", 1024 ** 3), ("PB", 1024 ** 4)):
        f = n / denom
        if f < 10.0:
            return "%.2f %s" % (f, unit)
        if f < 100.0:
            return "%d %s" % (round(f, 1), unit)
        if f < 1000.0:
            return "%d %s" % (round(f, 0), unit)
    return ">1TB ?!?"


## for i in range(51):
##     print(2**i-1, human_memory(2**i-1))
##     print(2**i, human_memory(2**i))
##     print(2**i+1, human_memory(2**i+1))


def index_command(parser, args):
    reporter = command_support.get_reporter(args.quiet)

    if args.title:
        title = args.title
    else:
        title = "MMPs from %r" % (args.fragment_filename,)

    fragment_filter = get_fragment_filter_from_args(parser, args)

    report_memory = False
    if args.memory:
        if psutil is None:
            sys.stderr.write("WARNING: Cannot report memory because the 'psutil' module is not available.\n")
        else:
            report_memory = True

    start_properties_memory = get_memory_use()
    if args.properties is None:
        selected_ids = properties = None
    else:
        try:
            properties_file = open_universal(args.properties)
        except IOError as err:
            parser.error("Cannot open --properties file: %s" % (err,))

        try:
            with properties_file:
                properties = properties_io.load_properties(properties_file, reporter)
        except ValueError as err:
            parser.error("Problem reading --properties file %r: %s"
                         % (args.properties, err))

        selected_ids = set(properties.get_ids())
    end_properties_memory = get_memory_use()

    if (args.out is None or args.out == "mmpdb") and args.output is None:
        # Use the filename based on the fragments filename
        fragment_filename = args.fragment_filename
        if fragment_filename is None:
            parser.error("The '--out mmpdb' format requires a filename when reading from stdin.")

        # replace the extension (if any) with ".mmpdb"
        args.output = os.path.splitext(fragment_filename)[0] + ".mmpdb"
        reporter.warning("No --output filename specified. Saving to %r." % (args.output,))

    # reporter.report("Using fragment filters: %s" % (fragment_filter.get_args(),))

    fragment_io.suggest_faster_json(reporter)

    start_fragment_index_memory = get_memory_use()
    try:
        fragment_reader = fragment_io.read_fragment_records(args.fragment_filename)
    except fragment_types.FragmentFormatError as err:
        parser.error(str(err))

    with fragment_reader:
        with reporter.progress(fragment_reader, "Loaded fragment record") as report_fragment_reader:
            fragment_index = index_algorithm.load_fragment_index(report_fragment_reader, fragment_filter, selected_ids,
                                                                 args.norm_smi_file)

    index_options = fragment_filter.get_options()
    index_options.update(
        symmetric=args.symmetric,
        max_heavies_transf=args.max_heavies_transf,
        max_frac_trans=args.max_frac_trans,
        smallest_transformation_only=args.smallest_transformation_only
    )
    index_options = index_algorithm.IndexOptions(**index_options)

    start_mmp_memory = get_memory_use()
    environment_cache = index_algorithm.EnvironmentCache()

    # get the rule (fp, radius, smirks and count)
    RULESWITHENV_DICT, INDEX_TO_FINGERPRINT, INDEX_TO_ENVSMI = index_algorithm.findRulesAndEnvs(
        fragment_index, fragment_reader, index_options, environment_cache,
        max_radius=args.max_radius, reporter=reporter)

    # only possible at the moment with csv
    # write out to table
    with open(args.output, "w") as fwriter:
        fwriter.write("rule_smirks\tradius\tenvfp\tenv\tfrequency\n")
        for rule in RULESWITHENV_DICT:
            rule_fp = INDEX_TO_FINGERPRINT[rule[0]]
            rule_radius = rule[1]
            rule_smirks = rule[2]
            rule_env = INDEX_TO_ENVSMI[rule[3]]
            freq = RULESWITHENV_DICT[rule]
            fwriter.write(
                rule_smirks + "\t" + str(rule_radius) + "\t" + str(rule_fp) + "\t" + rule_env + "\t" + str(freq) + "\n")

cfd0621f4287d7e777f263441ddb21908e7f13f0
<mask token>

<mask token>


def get_fragment_filter_from_args(parser, args):
    min_variable_heavies = args.min_variable_heavies
    max_variable_heavies = args.max_variable_heavies
    min_variable_ratio = args.min_variable_ratio
    max_variable_ratio = args.max_variable_ratio
    if max_variable_heavies == 'none':
        max_variable_heavies = None
    if min_variable_ratio is not None and max_variable_ratio is not None:
        if min_variable_ratio > max_variable_ratio:
            parser.error(
                '--min-variable-ratio must not be larger than --max-variable-ratio'
                )
    if min_variable_heavies is not None and max_variable_heavies is not None:
        if min_variable_heavies > max_variable_heavies:
            parser.error(
                '--min-variable-heavies must not be larger than --max-variable-heavies'
                )
    filters = []
    if min_variable_heavies is not None:
        filters.append(index_algorithm.MinVariableHeaviesFilter(
            min_variable_heavies))
    if max_variable_heavies is not None:
        filters.append(index_algorithm.MaxVariableHeaviesFilter(
            max_variable_heavies))
    if min_variable_ratio is not None:
        filters.append(index_algorithm.MinVariableRatioFilter(
            min_variable_ratio))
    if max_variable_ratio is not None:
        filters.append(index_algorithm.MaxVariableRatioFilter(
            max_variable_ratio))
    if not filters:
        return index_algorithm.MultipleFilters([])
    elif len(filters) == 1:
        return filters[0]
    else:
        return index_algorithm.MultipleFilters(filters)


<mask token>


def human_memory(n):
    if n < 1024:
        return '%d B' % (n,)
    for unit, denom in (('KB', 1024), ('MB', 1024 ** 2), ('GB', 1024 ** 3),
        ('PB', 1024 ** 4)):
        f = n / denom
        if f < 10.0:
            return '%.2f %s' % (f, unit)
        if f < 100.0:
            return '%d %s' % (round(f, 1), unit)
        if f < 1000.0:
            return '%d %s' % (round(f, 0), unit)
    return '>1TB ?!?'


def index_command(parser, args):
    reporter = command_support.get_reporter(args.quiet)
    if args.title:
        title = args.title
    else:
        title = 'MMPs from %r' % (args.fragment_filename,)
    fragment_filter = get_fragment_filter_from_args(parser, args)
    report_memory = False
    if args.memory:
        if psutil is None:
            sys.stderr.write(
                """WARNING: Cannot report memory because the 'psutil' module is not available.
"""
                )
        else:
            report_memory = True
    start_properties_memory = get_memory_use()
    if args.properties is None:
        selected_ids = properties = None
    else:
        try:
            properties_file = open_universal(args.properties)
        except IOError as err:
            parser.error('Cannot open --properties file: %s' % (err,))
        try:
            with properties_file:
                properties = properties_io.load_properties(properties_file,
                    reporter)
        except ValueError as err:
            parser.error('Problem reading --properties file %r: %s' % (args
                .properties, err))
        selected_ids = set(properties.get_ids())
    end_properties_memory = get_memory_use()
    if (args.out is None or args.out == 'mmpdb') and args.output is None:
        fragment_filename = args.fragment_filename
        if fragment_filename is None:
            parser.error(
                "The '--out mmpdb' format requires a filename when reading from stdin."
                )
        args.output = os.path.splitext(fragment_filename)[0] + '.mmpdb'
        reporter.warning('No --output filename specified. Saving to %r.' %
            (args.output,))
    fragment_io.suggest_faster_json(reporter)
    start_fragment_index_memory = get_memory_use()
    try:
        fragment_reader = fragment_io.read_fragment_records(args.
            fragment_filename)
    except fragment_types.FragmentFormatError as err:
        parser.error(str(err))
    with fragment_reader:
        with reporter.progress(fragment_reader, 'Loaded fragment record'
            ) as report_fragment_reader:
            fragment_index = index_algorithm.load_fragment_index(
                report_fragment_reader, fragment_filter, selected_ids, args
                .norm_smi_file)
    index_options = fragment_filter.get_options()
    index_options.update(symmetric=args.symmetric, max_heavies_transf=args.
        max_heavies_transf, max_frac_trans=args.max_frac_trans,
        smallest_transformation_only=args.smallest_transformation_only)
    index_options = index_algorithm.IndexOptions(**index_options)
    start_mmp_memory = get_memory_use()
    environment_cache = index_algorithm.EnvironmentCache()
    RULESWITHENV_DICT, INDEX_TO_FINGERPRINT, INDEX_TO_ENVSMI = (index_algorithm
        .findRulesAndEnvs(fragment_index, fragment_reader, index_options,
        environment_cache, max_radius=args.max_radius, reporter=reporter))
    with open(args.output, 'w') as fwriter:
        fwriter.write('rule_smirks\tradius\tenvfp\tenv\tfrequency\n')
        for rule in RULESWITHENV_DICT:
            rule_fp = INDEX_TO_FINGERPRINT[rule[0]]
            rule_radius = rule[1]
            rule_smirks = rule[2]
            rule_env = INDEX_TO_ENVSMI[rule[3]]
            freq = RULESWITHENV_DICT[rule]
            fwriter.write(rule_smirks + '\t' + str(rule_radius) + '\t' +
                str(rule_fp) + '\t' + rule_env + '\t' + str(freq) + '\n')

<mask token>
try:
    import psutil
except ImportError:
    psutil = None


def get_fragment_filter_from_args(parser, args):
    min_variable_heavies = args.min_variable_heavies
    max_variable_heavies = args.max_variable_heavies
    min_variable_ratio = args.min_variable_ratio
    max_variable_ratio = args.max_variable_ratio
    if max_variable_heavies == 'none':
        max_variable_heavies = None
    if min_variable_ratio is not None and max_variable_ratio is not None:
        if min_variable_ratio > max_variable_ratio:
            parser.error(
                '--min-variable-ratio must not be larger than --max-variable-ratio'
                )
    if min_variable_heavies is not None and max_variable_heavies is not None:
        if min_variable_heavies > max_variable_heavies:
            parser.error(
                '--min-variable-heavies must not be larger than --max-variable-heavies'
                )
    filters = []
    if min_variable_heavies is not None:
        filters.append(index_algorithm.MinVariableHeaviesFilter(
            min_variable_heavies))
    if max_variable_heavies is not None:
        filters.append(index_algorithm.MaxVariableHeaviesFilter(
            max_variable_heavies))
    if min_variable_ratio is not None:
        filters.append(index_algorithm.MinVariableRatioFilter(
            min_variable_ratio))
    if max_variable_ratio is not None:
        filters.append(index_algorithm.MaxVariableRatioFilter(
            max_variable_ratio))
    if not filters:
        return index_algorithm.MultipleFilters([])
    elif len(filters) == 1:
        return filters[0]
    else:
        return index_algorithm.MultipleFilters(filters)


if psutil is None:

    def get_memory_use():
        return 0
else:
    _process = psutil.Process(os.getpid())

    def get_memory_use():
        info = _process.memory_info()
        return info.rss


def human_memory(n):
    if n < 1024:
        return '%d B' % (n,)
    for unit, denom in (('KB', 1024), ('MB', 1024 ** 2), ('GB', 1024 ** 3),
        ('PB', 1024 ** 4)):
        f = n / denom
        if f < 10.0:
            return '%.2f %s' % (f, unit)
        if f < 100.0:
            return '%d %s' % (round(f, 1), unit)
        if f < 1000.0:
            return '%d %s' % (round(f, 0), unit)
    return '>1TB ?!?'


def index_command(parser, args):
    reporter = command_support.get_reporter(args.quiet)
    if args.title:
        title = args.title
    else:
        title = 'MMPs from %r' % (args.fragment_filename,)
    fragment_filter = get_fragment_filter_from_args(parser, args)
    report_memory = False
    if args.memory:
        if psutil is None:
            sys.stderr.write(
                """WARNING: Cannot report memory because the 'psutil' module is not available.
"""
                )
        else:
            report_memory = True
    start_properties_memory = get_memory_use()
    if args.properties is None:
        selected_ids = properties = None
    else:
        try:
            properties_file = open_universal(args.properties)
        except IOError as err:
            parser.error('Cannot open --properties file: %s' % (err,))
        try:
            with properties_file:
                properties = properties_io.load_properties(properties_file,
                    reporter)
        except ValueError as err:
            parser.error('Problem reading --properties file %r: %s' % (args
                .properties, err))
        selected_ids = set(properties.get_ids())
    end_properties_memory = get_memory_use()
    if (args.out is None or args.out == 'mmpdb') and args.output is None:
        fragment_filename = args.fragment_filename
        if fragment_filename is None:
            parser.error(
                "The '--out mmpdb' format requires a filename when reading from stdin."
                )
        args.output = os.path.splitext(fragment_filename)[0] + '.mmpdb'
        reporter.warning('No --output filename specified. Saving to %r.' %
            (args.output,))
    fragment_io.suggest_faster_json(reporter)
    start_fragment_index_memory = get_memory_use()
    try:
        fragment_reader = fragment_io.read_fragment_records(args.
            fragment_filename)
    except fragment_types.FragmentFormatError as err:
        parser.error(str(err))
    with fragment_reader:
        with reporter.progress(fragment_reader, 'Loaded fragment record'
            ) as report_fragment_reader:
            fragment_index = index_algorithm.load_fragment_index(
                report_fragment_reader, fragment_filter, selected_ids, args
                .norm_smi_file)
    index_options = fragment_filter.get_options()
    index_options.update(symmetric=args.symmetric, max_heavies_transf=args.
        max_heavies_transf, max_frac_trans=args.max_frac_trans,
        smallest_transformation_only=args.smallest_transformation_only)
    index_options = index_algorithm.IndexOptions(**index_options)
    start_mmp_memory = get_memory_use()
    environment_cache = index_algorithm.EnvironmentCache()
    RULESWITHENV_DICT, INDEX_TO_FINGERPRINT, INDEX_TO_ENVSMI = (index_algorithm
        .findRulesAndEnvs(fragment_index, fragment_reader, index_options,
        environment_cache, max_radius=args.max_radius, reporter=reporter))
    with open(args.output, 'w') as fwriter:
        fwriter.write('rule_smirks\tradius\tenvfp\tenv\tfrequency\n')
        for rule in RULESWITHENV_DICT:
            rule_fp = INDEX_TO_FINGERPRINT[rule[0]]
            rule_radius = rule[1]
            rule_smirks = rule[2]
            rule_env = INDEX_TO_ENVSMI[rule[3]]
            freq = RULESWITHENV_DICT[rule]
            fwriter.write(rule_smirks + '\t' + str(rule_radius) + '\t' +
                str(rule_fp) + '\t' + rule_env + '\t' + str(freq) + '\n')

from __future__ import print_function, absolute_import, division
import sys
import os
from . import command_support
from . import index_algorithm
from . import fragment_io
from . import fragment_types
from . import properties_io
from ._compat import open_universal
try:
    import psutil
except ImportError:
    psutil = None


def get_fragment_filter_from_args(parser, args):
    min_variable_heavies = args.min_variable_heavies
    max_variable_heavies = args.max_variable_heavies
    min_variable_ratio = args.min_variable_ratio
    max_variable_ratio = args.max_variable_ratio
    if max_variable_heavies == 'none':
        max_variable_heavies = None
    if min_variable_ratio is not None and max_variable_ratio is not None:
        if min_variable_ratio > max_variable_ratio:
            parser.error(
                '--min-variable-ratio must not be larger than --max-variable-ratio'
                )
    if min_variable_heavies is not None and max_variable_heavies is not None:
        if min_variable_heavies > max_variable_heavies:
            parser.error(
                '--min-variable-heavies must not be larger than --max-variable-heavies'
                )
    filters = []
    if min_variable_heavies is not None:
        filters.append(index_algorithm.MinVariableHeaviesFilter(
            min_variable_heavies))
    if max_variable_heavies is not None:
        filters.append(index_algorithm.MaxVariableHeaviesFilter(
            max_variable_heavies))
    if min_variable_ratio is not None:
        filters.append(index_algorithm.MinVariableRatioFilter(
            min_variable_ratio))
    if max_variable_ratio is not None:
        filters.append(index_algorithm.MaxVariableRatioFilter(
            max_variable_ratio))
    if not filters:
        return index_algorithm.MultipleFilters([])
    elif len(filters) == 1:
        return filters[0]
    else:
        return index_algorithm.MultipleFilters(filters)


if psutil is None:

    def get_memory_use():
        return 0
else:
    _process = psutil.Process(os.getpid())

    def get_memory_use():
        info = _process.memory_info()
        return info.rss


def human_memory(n):
    if n < 1024:
        return '%d B' % (n,)
    for unit, denom in (('KB', 1024), ('MB', 1024 ** 2), ('GB', 1024 ** 3),
        ('PB', 1024 ** 4)):
        f = n / denom
        if f < 10.0:
            return '%.2f %s' % (f, unit)
        if f < 100.0:
            return '%d %s' % (round(f, 1), unit)
        if f < 1000.0:
            return '%d %s' % (round(f, 0), unit)
    return '>1TB ?!?'


def index_command(parser, args):
    reporter = command_support.get_reporter(args.quiet)
    if args.title:
        title = args.title
    else:
        title = 'MMPs from %r' % (args.fragment_filename,)
    fragment_filter = get_fragment_filter_from_args(parser, args)
    report_memory = False
    if args.memory:
        if psutil is None:
            sys.stderr.write(
                """WARNING: Cannot report memory because the 'psutil' module is not available.
"""
                )
        else:
            report_memory = True
    start_properties_memory = get_memory_use()
    if args.properties is None:
        selected_ids = properties = None
    else:
        try:
            properties_file = open_universal(args.properties)
        except IOError as err:
            parser.error('Cannot open --properties file: %s' % (err,))
        try:
            with properties_file:
                properties = properties_io.load_properties(properties_file,
                    reporter)
        except ValueError as err:
            parser.error('Problem reading --properties file %r: %s' % (args
                .properties, err))
        selected_ids = set(properties.get_ids())
    end_properties_memory = get_memory_use()
    if (args.out is None or args.out == 'mmpdb') and args.output is None:
        fragment_filename = args.fragment_filename
        if fragment_filename is None:
            parser.error(
                "The '--out mmpdb' format requires a filename when reading from stdin."
                )
        args.output = os.path.splitext(fragment_filename)[0] + '.mmpdb'
        reporter.warning('No --output filename specified. Saving to %r.' %
            (args.output,))
    fragment_io.suggest_faster_json(reporter)
    start_fragment_index_memory = get_memory_use()
    try:
        fragment_reader = fragment_io.read_fragment_records(args.
            fragment_filename)
    except fragment_types.FragmentFormatError as err:
        parser.error(str(err))
    with fragment_reader:
        with reporter.progress(fragment_reader, 'Loaded fragment record'
            ) as report_fragment_reader:
            fragment_index = index_algorithm.load_fragment_index(
                report_fragment_reader, fragment_filter, selected_ids, args
                .norm_smi_file)
    index_options = fragment_filter.get_options()
    index_options.update(symmetric=args.symmetric, max_heavies_transf=args.
        max_heavies_transf, max_frac_trans=args.max_frac_trans,
        smallest_transformation_only=args.smallest_transformation_only)
    index_options = index_algorithm.IndexOptions(**index_options)
    start_mmp_memory = get_memory_use()
    environment_cache = index_algorithm.EnvironmentCache()
    RULESWITHENV_DICT, INDEX_TO_FINGERPRINT, INDEX_TO_ENVSMI = (index_algorithm
        .findRulesAndEnvs(fragment_index, fragment_reader, index_options,
        environment_cache, max_radius=args.max_radius, reporter=reporter))
    with open(args.output, 'w') as fwriter:
        fwriter.write('rule_smirks\tradius\tenvfp\tenv\tfrequency\n')
        for rule in RULESWITHENV_DICT:
            rule_fp = INDEX_TO_FINGERPRINT[rule[0]]
            rule_radius = rule[1]
            rule_smirks = rule[2]
            rule_env = INDEX_TO_ENVSMI[rule[3]]
            freq = RULESWITHENV_DICT[rule]
            fwriter.write(rule_smirks + '\t' + str(rule_radius) + '\t' +
                str(rule_fp) + '\t' + rule_env + '\t' + str(freq) + '\n')

# mmpdb - matched molecular pair database generation and analysis
#
# Copyright (c) 2015-2017, F. Hoffmann-La Roche Ltd.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#    * Redistributions of source code must retain the above copyright
#      notice, this list of conditions and the following disclaimer.
#    * Redistributions in binary form must reproduce the above
#      copyright notice, this list of conditions and the following
#      disclaimer in the documentation and/or other materials provided
#      with the distribution.
#    * Neither the name of F. Hoffmann-La Roche Ltd. nor the names of
#      its contributors may be used to endorse or promote products
#      derived from this software without specific prior written
#      permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from __future__ import print_function, absolute_import, division

import sys
import os

from . import command_support
from . import index_algorithm
from . import fragment_io
from . import fragment_types
from . import properties_io
from ._compat import open_universal

try:
    import psutil
except ImportError:
    psutil = None


def get_fragment_filter_from_args(parser, args):
    min_variable_heavies = args.min_variable_heavies
    max_variable_heavies = args.max_variable_heavies
    min_variable_ratio = args.min_variable_ratio
    max_variable_ratio = args.max_variable_ratio

    if max_variable_heavies == "none":  # Push this into the argument parser?
        max_variable_heavies = None

    if min_variable_ratio is not None and max_variable_ratio is not None:
        if min_variable_ratio > max_variable_ratio:
            parser.error("--min-variable-ratio must not be larger than --max-variable-ratio")

    if min_variable_heavies is not None and max_variable_heavies is not None:
        if min_variable_heavies > max_variable_heavies:
            parser.error("--min-variable-heavies must not be larger than --max-variable-heavies")

    filters = []
    if min_variable_heavies is not None:
        filters.append(index_algorithm.MinVariableHeaviesFilter(min_variable_heavies))
    if max_variable_heavies is not None:
        filters.append(index_algorithm.MaxVariableHeaviesFilter(max_variable_heavies))
    if min_variable_ratio is not None:
        filters.append(index_algorithm.MinVariableRatioFilter(min_variable_ratio))
    if max_variable_ratio is not None:
        filters.append(index_algorithm.MaxVariableRatioFilter(max_variable_ratio))

    if not filters:
        # It's easier to have 0 filters than to make a special do-nothing filter.
        return index_algorithm.MultipleFilters([])
    elif len(filters) == 1:
        return filters[0]
    else:
        return index_algorithm.MultipleFilters(filters)


if psutil is None:
    def get_memory_use():
        return 0
else:
    _process = psutil.Process(os.getpid())


    def get_memory_use():
        info = _process.memory_info()
        return info.rss  # or info.vms?


def human_memory(n):
    if n < 1024:
        return "%d B" % (n,)
    for unit, denom in (("KB", 1024), ("MB", 1024 ** 2),
                        ("GB", 1024 ** 3), ("PB", 1024 ** 4)):
        f = n / denom
        if f < 10.0:
            return "%.2f %s" % (f, unit)
        if f < 100.0:
            return "%d %s" % (round(f, 1), unit)
        if f < 1000.0:
            return "%d %s" % (round(f, 0), unit)
    return ">1TB ?!?"


## for i in range(51):
##     print(2**i-1, human_memory(2**i-1))
##     print(2**i, human_memory(2**i))
##     print(2**i+1, human_memory(2**i+1))


def index_command(parser, args):
    reporter = command_support.get_reporter(args.quiet)

    if args.title:
        title = args.title
    else:
        title = "MMPs from %r" % (args.fragment_filename,)

    fragment_filter = get_fragment_filter_from_args(parser, args)

    report_memory = False
    if args.memory:
        if psutil is None:
            sys.stderr.write("WARNING: Cannot report memory because the 'psutil' module is not available.\n")
        else:
            report_memory = True

    start_properties_memory = get_memory_use()
    if args.properties is None:
        selected_ids = properties = None
    else:
        try:
            properties_file = open_universal(args.properties)
        except IOError as err:
            parser.error("Cannot open --properties file: %s" % (err,))

        try:
            with properties_file:
                properties = properties_io.load_properties(properties_file, reporter)
        except ValueError as err:
            parser.error("Problem reading --properties file %r: %s"
                         % (args.properties, err))

        selected_ids = set(properties.get_ids())
    end_properties_memory = get_memory_use()

    if (args.out is None or args.out == "mmpdb") and args.output is None:
        # Use the filename based on the fragments filename
        fragment_filename = args.fragment_filename
        if fragment_filename is None:
            parser.error("The '--out mmpdb' format requires a filename when reading from stdin.")

        # replace the extension (if any) with ".mmpdb"
        args.output = os.path.splitext(fragment_filename)[0] + ".mmpdb"
        reporter.warning("No --output filename specified. Saving to %r." % (args.output,))

    # reporter.report("Using fragment filters: %s" % (fragment_filter.get_args(),))

    fragment_io.suggest_faster_json(reporter)

    start_fragment_index_memory = get_memory_use()
    try:
        fragment_reader = fragment_io.read_fragment_records(args.fragment_filename)
    except fragment_types.FragmentFormatError as err:
        parser.error(str(err))

    with fragment_reader:
        with reporter.progress(fragment_reader, "Loaded fragment record") as report_fragment_reader:
            fragment_index = index_algorithm.load_fragment_index(report_fragment_reader, fragment_filter, selected_ids,
                                                                 args.norm_smi_file)

    index_options = fragment_filter.get_options()
    index_options.update(
        symmetric=args.symmetric,
        max_heavies_transf=args.max_heavies_transf,
        max_frac_trans=args.max_frac_trans,
        smallest_transformation_only=args.smallest_transformation_only
    )
    index_options = index_algorithm.IndexOptions(**index_options)

    start_mmp_memory = get_memory_use()
    environment_cache = index_algorithm.EnvironmentCache()

    # get the rule (fp, radius, smirks and count)
    RULESWITHENV_DICT, INDEX_TO_FINGERPRINT, INDEX_TO_ENVSMI = index_algorithm.findRulesAndEnvs(
        fragment_index, fragment_reader, index_options, environment_cache,
        max_radius=args.max_radius, reporter=reporter)

    # only possible at the moment with csv
    # write out to table
    with open(args.output, "w") as fwriter:
        fwriter.write("rule_smirks\tradius\tenvfp\tenv\tfrequency\n")
        for rule in RULESWITHENV_DICT:
            rule_fp = INDEX_TO_FINGERPRINT[rule[0]]
            rule_radius = rule[1]
            rule_smirks = rule[2]
            rule_env = INDEX_TO_ENVSMI[rule[3]]
            freq = RULESWITHENV_DICT[rule]
            fwriter.write(
                rule_smirks + "\t" + str(rule_radius) + "\t" + str(rule_fp) + "\t" + rule_env + "\t" + str(freq) + "\n")
