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

r"""
Read hyb / vienna files and analyze the fold information in the contained hybrid sequences.

Analysis Types:

    =============================== ==============================================================
    :ref:`Energy <EnergyAnalysis>`  Analysis of values of predicted intra-hybrid folding energy
    :ref:`Type <TypeAnalysis>`      Analysis of segment types
    :ref:`miRNA <MirnaAnalysis>`    Analysis of miRNA segments distributions
    :ref:`Target <TargetAnalysis>`  Analysis of mirna target segment names and types
    :ref:`Fold <FoldAnalysis>`      Analysis of folding data included in the analyzed hyb_records
    =============================== ==============================================================

This utility reads in one or more files in hyb-format
(see the :ref:`Hybkit Hyb File Specification`)
along with a corresponding predicted RNA secondary structure file in the
"Vienna" (:ref:`Vienna Format <vienna_file_format>`) or
"CT" (:ref:`CT_Format <ct_file_format>`) formats,
and analyzes hybrid secondary structure properties.

For full information on the different analysis types, see the :ref:`Analyses <Analyses>`
section of the hybkit documentation.

Example system calls:
    ::

        $ hyb_analyze -a fold -i my_file_1.hyb -f my_file_1.vienna

        $ hyb_analyze -a fold -i my_file_2.hyb -f my_file_2.ct \\
                    --make_plots False
"""

import argparse
import contextlib
import os
from typing import List, Optional, Union

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

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

# ----- Begin Custom Types -----
AnalysisTypesArg = Union[str, List[str]]

# Create Command-line Argument Parser
def make_parser() -> argparse.ArgumentParser:
    """Create the command-line argument parser for the hyb_analyze script."""
    parser_components = [
        hybkit.util.in_hybs_parser,
        hybkit.util.in_folds_parser,
        hybkit.util.out_basenames_parser,
        hybkit.util.cmb_out_opts_parser,
        hybkit.util.hyb_analyze_parser,
        hybkit.util.all_analyze_parser,
        hybkit.util.gen_opts_parser,
        hybkit.util.cmb_hyb_fold_class_settings_parser,
        hybkit.util.analysis_parser,
    ]

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

    return script_parser


# Define main script function.
def hyb_analyze(
        in_hyb_files: List[str],
        analysis_types: AnalysisTypesArg,
        in_fold_files: Optional[List[str]] = None,
        out_dir: str = '.',
        out_suffix: Optional[str] = None,
        out_basenames: Optional[List[str]] = None,
        analysis_name: Optional[str] = None,
        make_plots: bool = False,
        verbose: bool = False,
        silent: bool = False,
        ) -> None:
    """Perform main script function."""
    if not silent:
        print('\nPerforming Analysis of Hyb Files...')

    if verbose:
        print('\nPerforming Analysis Types: ' + ', '.join(analysis_types))
        if analysis_name:
            print('Analysis Title: "%s"' % analysis_name)
        print()

    if out_suffix is None:
        if len(analysis_types) > 1:
            out_suffix = '_multi-analysis'
        else:
            out_suffix = '_' + analysis_types[0]
    if verbose:
        print('Using Out Suffix: "%s"' % out_suffix)

    combined_analysis = hybkit.analysis.Analysis(analysis_types=analysis_types, name=analysis_name)

    # Start Setup Input Files
    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_basenames is not None:
            out_basename = out_basenames[i]
        else:
            out_basename = hybkit.util.make_out_file_name(
                in_hyb_file,
                name_suffix=out_suffix,
                in_suffix='.hyb',
                out_suffix='',
                out_dir=out_dir,
                seg_sep='_',
            )

        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)
            in_fold_args = (in_fold_file, 'r')
        else:
            in_fold_class = contextlib.nullcontext
            in_fold_args = ()

        if verbose:
            print('Analyzing Files:')
            print('    Input Hyb:   ' + in_hyb_file)
            if in_fold_files:
                print('    Input Fold:  ' + in_fold_file)
            print('    Output Base: ' + out_basename)

        # Start Record Iteration
        with hybkit.HybFile(in_hyb_file, 'r') as in_hyb, \
             in_fold_class(*in_fold_args) as in_fold:
            if in_fold_file is None:
                record_iter = in_hyb
            else:
                record_iter = hybkit.HybFoldIter(in_hyb, in_fold, combine=True)

            for _, hyb_record in enumerate(record_iter, start=1):
                combined_analysis.add_hyb_record(hyb_record)

        combined_analysis.write_analysis_results_special(out_basename=out_basename)
        if make_plots:
            combined_analysis.plot_analysis_results(out_basename=out_basename)

    if verbose:
        if hasattr(record_iter, 'print_report'):
            print('\nHybFoldIter Report:\n')
            record_iter.print_report()
        print('\nAnalysis 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)
    hyb_analyze(
        in_hyb_files=args.in_hyb,
        in_fold_files=args.in_fold,
        analysis_types=args.analysis_types,
        out_dir=args.out_dir,
        out_suffix=args.out_suffix,
        out_basenames=args.out_basename,
        analysis_name=args.analysis_name,
        make_plots=args.make_plots,
        verbose=args.verbose,
        silent=args.silent,
    )
