#!/usr/bin/env python3
# Daniel B. Stribling
# Renne Lab, University of Florida
# Hybkit Project : http://www.github.com/RenneLab/hybkit

r"""
Filter hyb (and corresponding fold) files to meet (or exclude) specific criteria.

This script takes one or more filter and/or exclusion criteria
and outputs only those records matching (/excluding) those criteria.

The filter criteria and options are based on the options provided by the
:func:`hybkit.HybRecord.prop` method of the Hybkit API. For more information see
the full documentation for the :class:`~hybkit.HybRecord` class.

Example System Calls:
    ::

        hyb_filter -i my_file_1.hyb --filter has_seg_types
        # Outputs records that have completed a segtype analysis

        hyb_filter -i my_file_1.hyb -f my_file_1.vienna \\
            --include seg_type mRNA
        # Outputs hyb and fold records where hyb record has either segtype of mRNA

        hyb_filter -i my_file_1.hyb --exclude seg_type mRNA
        # Outputs records without either segtype of mRNA

        hyb_filter -i my_file_1.hyb --include seg1_type mRNA
        # Outputs records with only the first / 5p segtype of mRNA

        hyb_filter -i my_file_1.hyb my_file_2.hyb -f my_file_1.vienna my_file_2.vienna \\
            --include seg_type_contains RNA
        # Outputs all records with a segtype that includes
        #   the string "RNA" (case-sensitive)

        hyb_filter -i my_file_1.hyb --filter seg_contains kshv
        # Outputs records where either segment identifier contains the
        #   the string: "kshv" (case-sensitive)

Multiple filtering options can be used together.
The ``-m`` / ``--filter_mode`` argument determines whether
"any" (DEFAULT) or "all" filters are required to be true for inclusion.
Note: Matching any exclusion criteria results in exclusion of the record.

Example System Calls (match ALL criteria):
    ::

        hyb_filter -i my_file_1.hyb -f my_file_1.vienna \\
                    --filter seg_contains kshv \\
                    --filter_2 seg_type miRNA
        # Outputs records with either reference sequence identifier containing "kshv"
        #   and with either segment having an assigned segtype of miRNA

Example System Calls (match ANY criteria):
    ::

        hyb_filter -i my_file_1.hyb --filter_mode any \\
                    --filter seg_type miRNA \\
                    --filter_2 seg_type lncRNA
        # Outputs records containing either segment type matching
        #   either "miRNA" or "lncRNA" (case-sensitive)
"""

import argparse
import contextlib
import os
import sys
from typing import List, Literal, Optional, Tuple

import hybkit
from hybkit.__about__ import (
    __author__,
    __contact__,
    __credits__,
    __date__,
    __deprecated__,
    __email__,
    __license__,
    __maintainer__,
    __status__,
    __version__,
)

# ----- Linting Directives:
# ruff: noqa: F401 SLF001

# Create Command-line Argument Parser
def make_parser() -> argparse.ArgumentParser:
    """Create and return the argparse.ArgumentParser for the hyb_filter script."""
    parser_components = [
        hybkit.util.cmb_hyb_fold_io_parser,
        hybkit.util.cmb_out_opts_parser,
        hybkit.util.hyb_filter_parser,
        hybkit.util.record_manip_parser,
        hybkit.util.gen_opts_parser,
        hybkit.util.cmb_hyb_fold_class_settings_parser,
    ]

    script_parser = argparse.ArgumentParser(
        parents=parser_components,
        prog='hyb_filter',
        description=hybkit.util.get_argparse_doc(__doc__),
        epilog=hybkit.util.output_description,
        formatter_class=hybkit.util._HybkitFormatter,
        allow_abbrev=False,
    )

    script_parser.set_defaults(out_suffix=hybkit.settings._FILTER_OUT_SUFFIX)

    return script_parser


# Define main script function.
def hyb_filter(
        in_hyb_files: List[str],
        filter_params: List[Tuple[str, str]],
        exclude_params: List[Tuple[str, str]],
        in_fold_files: Optional[List[str]] = None,
        filter_mode: Literal['all', 'any'] = 'all',
        out_dir: str = '.',
        out_suffix: str = hybkit.settings._FILTER_OUT_SUFFIX,
        out_hyb_files: Optional[List[str]] = None,
        out_fold_files: Optional[List[str]] = None,
        set_dataset: Optional[str] = None,
        skip_dup_id_before: bool = False,
        skip_dup_id_after: bool = False,
        verbose: bool = False,
        silent: bool = False,
        ) -> None:
    """Perform main script function."""
    if not silent:
        print('\nFiltering Hyb Files...')

    if verbose:
        if filter_params:
            print('Records Must match %s Filter Parameters:' % filter_mode.upper())
            for p_set in filter_params:
                print('   ', ' '.join([x for x in p_set if x is not None]))

        if exclude_params:
            print('Using Exclusion Filter Parameters:')
            print('    ' + '\n    '.join(' '.join(p) for p in exclude_params) + '\n')

    allowed_modes = {'all', 'any'}
    if filter_mode not in allowed_modes:
        message = 'Unrecognized filtering mode: %s\n' % filter_mode
        message += 'Options are: %s' % str(allowed_modes)

    if verbose:
        print('Using Out Suffix: "%s"' % out_suffix)

    if in_fold_files:
        file_iter = zip(in_hyb_files, in_fold_files)
    else:
        file_iter = in_hyb_files

    for i, use_files in enumerate(file_iter):
        if in_fold_files:
            in_hyb_file, in_fold_file = use_files
        else:
            in_hyb_file, in_fold_file = use_files, None
        file_basename = os.path.basename(in_hyb_file)
        file_label = file_basename.replace('.hyb', '')

        if out_hyb_files is not None:
            out_hyb_file = out_hyb_files[i]
        else:
            out_hyb_file = hybkit.util.make_out_file_name(
                in_hyb_file,
                name_suffix=out_suffix,
                in_suffix='.hyb',
                out_suffix='.hyb',
                out_dir=out_dir,
                seg_sep='_',
            )

        if in_fold_file is None:
            out_fold_file = None
        elif out_fold_files is not None:
            out_fold_file = out_fold_files[i]
        else:
            out_fold_file = out_hyb_file.replace('.hyb', '.vienna')

        if in_fold_file is not None:
            if any(in_fold_file.endswith(s) for s in hybkit.settings.VIENNA_SUFFIXES):
                in_fold_class = hybkit.ViennaFile
            elif any(in_fold_file.endswith(s) for s in hybkit.settings.CT_SUFFIXES):
                in_fold_class = hybkit.CtFile
            else:
                raise ValueError('Unrecognized fold file type: %s' % in_fold_file)
            out_fold_class = hybkit.ViennaFile
            in_fold_args = (in_fold_file, 'r')
            out_fold_args = (out_fold_file, 'w')
        else:
            in_fold_class = contextlib.nullcontext
            out_fold_class = contextlib.nullcontext
            in_fold_args = ()
            out_fold_args = ()

        if verbose:
            print('Filtering File:')
            print('    Input Hyb:   ' + in_hyb_file)
            if in_fold_files:
                print('    Input Fold:  ' + in_fold_file)
            print('    Output Hyb:  ' + out_hyb_file)
            if in_fold_files:
                print('    Output Fold: ' + out_fold_file)

        include_count = 0
        exclude_count = 0

        with hybkit.HybFile(in_hyb_file, 'r') as in_hyb, \
             hybkit.HybFile(out_hyb_file, 'w') as out_hyb, \
             in_fold_class(*in_fold_args) as in_fold, \
             out_fold_class(*out_fold_args) as out_fold:
            if in_fold_file is None:
                record_iter = in_hyb
            else:
                record_iter = hybkit.HybFoldIter(in_hyb, in_fold, combine=True)

            last_record_id = None

            for _i, hyb_record in enumerate(record_iter, start=1):
                if skip_dup_id_before:
                    if hyb_record.id == last_record_id:
                        continue
                    last_record_id = hyb_record.id

                if filter_mode == 'all':
                    use_record = True
                    for prop_type, prop_compare in filter_params:
                        if not hyb_record.has_prop(prop_type, prop_compare):
                            use_record = False
                            break

                elif filter_mode == 'any':
                    use_record = False
                    for prop_type, prop_compare in filter_params:
                        if hyb_record.has_prop(prop_type, prop_compare):
                            use_record = True
                            break

                if use_record:
                    for prop_type, prop_compare in exclude_params:
                        if hyb_record.has_prop(prop_type, prop_compare):
                            use_record = False
                            break

                if skip_dup_id_after and use_record:
                    if hyb_record.id == last_record_id:
                        use_record = False
                    else:
                        last_record_id = hyb_record.id

                if use_record:
                    if set_dataset:
                        hyb_record.set_flag('dataset', file_label)

                    out_hyb.write_record(hyb_record)
                    include_count += 1
                    if in_fold_file is not None:
                        out_fold.write_record(hyb_record.fold_record)

                else:
                    exclude_count += 1

        total_count = include_count + exclude_count
        if verbose:
            print('    Complete. %i Total,  ' % total_count
                  + '%i Included,  %i Excluded\n' % (include_count, exclude_count))

    if verbose:
        if hasattr(record_iter, 'print_report'):
            print('\nHybFoldIter Report:\n')
            record_iter.print_report()
        print('\nFiltering Complete.\n')


# Execute the script function
if __name__ == '__main__':
    script_parser = make_parser()
    args = script_parser.parse_args()
    hybkit.util.validate_args(args, script_parser)
    hybkit.util.set_settings_from_namespace(args, verbose=args.verbose)
    filter_params = []
    for param_set in [getattr(args, key) for key in ['filter', 'filter_2', 'filter_3']]:
        if param_set is not None and bool(param_set):
            if len(param_set) > 2:  # noqa: PLR2004
                print('Invalid Filtering Parameter Set Encountered: %s ' % str(param_set))
                print('Parameter sets should have 1 or 2 arguments.\n')
                sys.exit()
            elif len(param_set) == 1:
                param_set.append(None)
            filter_params.append(param_set)

    exclude_params = []
    for param_set in [getattr(args, key) for key in ['exclude', 'exclude_2', 'exclude_3']]:
        if param_set is not None and bool(param_set):
            if len(param_set) > 2:  # noqa: PLR2004
                print('Invalid Exclusion Filtering Parameter Set Encountered: '
                      + '%s ' % str(param_set))
                print('Parameter sets should have 1 or 2 arguments.\n')
                sys.exit()
            elif len(param_set) == 1:
                param_set.append(None)
            exclude_params.append(param_set)

    hyb_filter(
        in_hyb_files=args.in_hyb,
        in_fold_files=args.in_fold,
        filter_params=filter_params,
        exclude_params=exclude_params,
        out_dir=args.out_dir,
        out_suffix=args.out_suffix,
        out_hyb_files=args.out_hyb,
        out_fold_files=args.out_fold,
        set_dataset=args.set_dataset,
        skip_dup_id_before=args.skip_dup_id_before,
        skip_dup_id_after=args.skip_dup_id_after,
        verbose=args.verbose,
        silent=args.silent,
    )
