"""Parsers module provides all implementations of command line interfaces.

Each command should implment a `register_` function and `run_` function. The
register function should take a parser and add the appropriate entries for the
command. The run function should accept the parser.parse_args() object and
execute the command.  run commands should contain minimal logic. If run
functions become too complex consider moving logic into an appropriate betta
module. Imports required for a particular command should be made inside of the
run commands to avoid loading all modules when a user simply wants the help
string.
"""

import os
import sys
import atexit
import argparse
from pathlib import Path
from shutil import rmtree

from remora import constants
from remora import log, RemoraError

LOGGER = log.get_logger()

class SubcommandHelpFormatter(
    argparse.RawDescriptionHelpFormatter,
    argparse.ArgumentDefaultsHelpFormatter,
):
    """Helper function to prettier print subcommand help. This removes some
    extra lines of output when a final command parser is not selected.
    """

    def _format_action(self, action):
        parts = super(SubcommandHelpFormatter, self)._format_action(action)
        if action.nargs == argparse.PARSER:
            parts = "\n".join(parts.split("\n")[1:])
        return parts

##################
# remora dataset #
##################

def register_dataset(parser):
    subparser = parser.add_parser(
        "dataset",
        description="Remora dataset operations",
        help="Create or perform operations on a Remora dataset",
        formatter_class=SubcommandHelpFormatter,
    )
    ssubparser = subparser.add_subparsers(title="dataset commands")
    #  Since `dataset` has several sub-commands, print help as default
    subparser.set_defaults(func=lambda x: subparser.print_help())
    #  Register dataset sub commands
    register_dataset_prepare(ssubparser)
    register_dataset_split(ssubparser)
    register_dataset_merge(ssubparser)
    register_dataset_inspect(ssubparser)

def register_dataset_prepare(parser):
    subparser = parser.add_parser(
        "prepare",
        description="Prepare Remora training dataset",
        help="Prepare Remora training dataset.",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "pod5",
        help="POD5 file corresponding to bam file.",
    )
    subparser.add_argument(
        "bam",
        help="BAM file containing mv tags.",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument(
        "--output-remora-training-file",
        default="remora_training_dataset.npz",
        help="Output Remora training dataset file.",
    )
    out_grp.add_argument(
        "--log-filename",
        help="Log filename. Default: Don't output log file.",
    )

    data_grp = subparser.add_argument_group("Data Arguments")
    data_grp.add_argument(
        "--motif",
        nargs=2,
        action="append",
        metavar=("MOTIF", "FOCUS_POSITION"),
        help="""Extract training chunks centered on a defined motif. Argument
        takes 2 values representing 1) sequence motif and 2) focus position
        within the motif. For example to restrict to CpG sites use "
        "--motif CG 0". Default: Any context ("N 0")""",
    )
    data_grp.add_argument(
        "--focus-reference-positions",
        help="""BED file containing reference positions around which to extract
        training chunks.""",
    )
    data_grp.add_argument(
        "--chunk-context",
        default=constants.DEFAULT_CHUNK_CONTEXT,
        type=int,
        nargs=2,
        metavar=("NUM_BEFORE", "NUM_AFTER"),
        help="""Number of context signal points to select around the central
        position.""",
    )
    data_grp.add_argument(
        "--min-samples-per-base",
        type=int,
        default=constants.DEFAULT_MIN_SAMPLES_PER_BASE,
        help="""Minimum number of samples per base. This sets the size of the
        ragged arrays of chunk sequences.""",
    )
    data_grp.add_argument(
        "--kmer-context-bases",
        nargs=2,
        default=constants.DEFAULT_KMER_CONTEXT_BASES,
        type=int,
        metavar=("BASES_BEFORE", "BASES_AFTER"),
        help="""Definition of k-mer (derived from the reference) passed into
        the model along with each signal position.""",
    )
    data_grp.add_argument(
        "--max-chunks-per-read",
        type=int,
        default=15,
        help="Maxiumum number of chunks to extract from a single read.",
    )
    data_grp.add_argument(
        "--base-start-justify",
        action="store_true",
        help="""Justify extracted chunk against the start of the base of
        interest. Default justifies chunk to middle of signal of the base
        of interest.""",
    )
    data_grp.add_argument(
        "--offset",
        default=0,
        type=int,
        help="Offset selected chunk position by a number of bases.",
    )
    data_grp.add_argument(
        "--num-reads",
        type=int,
        help="Number of reads.",
    )
    data_grp.add_argument(
        "--basecall-anchor",
        action="store_true",
        help="""Make dataset from basecall sequence instead of aligned
        reference sequence""",
    )
    data_grp.add_argument(
        "--reverse-signal",
        action="store_true",
        help="""Is nanopore signal 3' to 5' orientation? Primarily for direct
        RNA""",
    )

    refine_grp = subparser.add_argument_group("Signal Mapping Refine Arguments")
    refine_grp.add_argument(
        "--refine-kmer-level-table",
        help="""Tab-delimited file containing no header and two fields:
        1. string k-mer sequence and 2. float expected normalized level.
        All k-mers must be the same length and all combinations of the bases
        'ACGT' must be present in the file.""",
    )
    refine_grp.add_argument(
        "--refine-rough-rescale",
        action="store_true",
        help="""Apply a rough rescaling using quantiles of signal+move table
        and levels.""",
    )
    refine_grp.add_argument(
        "--refine-scale-iters",
        default=-1,
        type=int,
        help="""Number of iterations of signal mapping refinement and signal
        re-scaling to perform. Set to 0 in order to perform signal mapping
        refinement, but skip re-scaling. Set to -1 (default) to skip signal
        mapping (potentially using levels for rough rescaling).""",
    )
    refine_grp.add_argument(
        "--refine-half-bandwidth",
        default=constants.DEFAULT_REFINE_HBW,
        type=int,
        help="""Half bandwidth around signal mapping over which to search for
        "new path.""",
    )
    refine_grp.add_argument(
        "--refine-algo",
        default=constants.DEFAULT_REFINE_ALGO,
        choices=constants.REFINE_ALGOS,
        help="Refinement algorithm to apply (if kmer level table is provided).",
    )
    refine_grp.add_argument(
        "--refine-short-dwell-parameters",
        default=constants.DEFAULT_REFINE_SHORT_DWELL_PARAMS,
        type=float,
        nargs=3,
        metavar=("TARGET", "LIMIT", "WEIGHT"),
        help="""Short dwell penalty refiner parameters. Dwells shorter than
        LIMIT will be penalized a value of `WEIGHT * (dwell - TARGET)^2`.""",
    )

    label_grp = subparser.add_argument_group("Label Arguments")
    label_grp.add_argument(
        "--mod-base",
        nargs=2,
        metavar=("SINGLE_LETTER_CODE", "MOD_BASE"),
        help="Modified base information. Example: `--mod-base m 5mC`",
    )
    label_grp.add_argument(
        "--mod-base-control",
        action="store_true",
        help="Is this a modified bases control sample?",
    )
    label_grp.add_argument(
        "--base-pred",
        action="store_true",
        help="Train to predict bases (SNPs) and not mods.",
    )

    comp_grp = subparser.add_argument_group("Compute Arguments")
    comp_grp.add_argument(
        "--num-extract-alignment-workers",
        type=int,
        default=1,
        help="Number of signal extraction workers.",
    )
    comp_grp.add_argument(
        "--num-extract-chunks-workers",
        type=int,
        default=1,
        help="""Number of chunk extraction workers. If performing signal
        refinement this should be increased.""",
    )

    subparser.set_defaults(func=run_dataset_prepare)

def run_dataset_prepare(args):
    from remora.util import Motif
    from remora.io import parse_bed
    from remora.refine_signal_map import SigMapRefiner
    from remora.prepare_train_data import extract_chunk_dataset

    if args.log_filename is not None:
        log.init_logger(args.log_filename)
    if args.mod_base is None and not args.mod_base_control:
        LOGGER.error("Must specify either --mod-base or --mod-base-control")
        sys.exit(1)
    motifs = [("N", 0)] if args.motif is None else args.motif
    motifs = [Motif(*mo) for mo in motifs]
    focus_ref_pos = (
        None
        if args.focus_reference_positions is None
        else parse_bed(args.focus_reference_positions)
    )

    sig_map_refiner = SigMapRefiner(
        kmer_model_filename=args.refine_kmer_level_table,
        do_rough_rescale=args.refine_rough_rescale,
        scale_iters=args.refine_scale_iters,
        algo=args.refine_algo,
        half_bandwidth=args.refine_half_bandwidth,
        sd_params=args.refine_short_dwell_parameters,
        do_fix_guage=True,
    )
    extract_chunk_dataset(
        args.bam,
        args.pod5,
        args.output_remora_training_file,
        args.mod_base,
        args.mod_base_control,
        motifs,
        focus_ref_pos,
        args.chunk_context,
        args.min_samples_per_base,
        args.max_chunks_per_read,
        sig_map_refiner,
        args.base_pred,
        args.kmer_context_bases,
        args.base_start_justify,
        args.offset,
        args.num_reads,
        args.num_extract_alignment_workers,
        args.num_extract_chunks_workers,
        basecall_anchor=args.basecall_anchor,
        rev_sig=args.reverse_signal,
    )

def register_dataset_split(parser):
    subparser = parser.add_parser(
        "split",
        description="Split Remora dataset",
        help="Split Remora dataset",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "input_remora_dataset",
        help="Remora training dataset to be split",
    )
    subparser.add_argument(
        "--output-basename",
        default="split_remora_dataset",
        help="Basename for output datasets.",
    )
    subparser.add_argument(
        "--val-prop",
        type=float,
        help="""The proportion of data to be split into validation set, where
        val-prop in [0,0.5).""",
    )
    subparser.add_argument(
        "--val-num",
        type=int,
        help="Number of validation chunks to select.",
    )
    subparser.add_argument(
        "--unstratified",
        action="store_true",
        help="""For --val-prop split, perform unstratified splitting. Default
        will perform split stratified over labels.""",
    )
    subparser.add_argument(
        "--by-label",
        action="store_true",
        help="Split dataset into one dataset for each unique label.",
    )
    subparser.set_defaults(func=run_dataset_split)

def run_dataset_split(args):
    from remora.data_chunks import RemoraDataset

    dataset = RemoraDataset.load_from_file(
        args.input_remora_dataset,
        shuffle_on_iter=False,
        drop_last=False,
    )
    LOGGER.info(f"Loaded dataset summary:\n{dataset.summary}")

    if args.by_label:
        for label, label_dataset in dataset.split_by_label():
            label_dataset.save(f"{args.output_basename}.{label}.npz")
            LOGGER.info(
                f"Wrote {label_dataset.nchunks} chunks to "
                f"{args.output_basename}.{label}.npz"
            )
    else:
        trn_set, val_set = dataset.split_data(
            val_prop=args.val_prop,
            val_num=args.val_num,
            stratified=not args.unstratified,
        )
        LOGGER.info(
            f"Train set label distribution: {trn_set.get_label_counts()}"
        )
        LOGGER.info(f"Val set label distribution: {val_set.get_label_counts()}")
        trn_set.save(f"{args.output_basename}.split_train.npz")
        val_set.save(f"{args.output_basename}.split_val.npz")

def register_dataset_inspect(parser):
    subparser = parser.add_parser(
        "inspect",
        description="Inspect Remora dataset",
        help="Inspect Remora dataset",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "remora_dataset_path",
        help="Remora training dataset",
    )
    subparser.set_defaults(func=run_dataset_inspect)

def run_dataset_inspect(args):
    from remora.data_chunks import RemoraDataset

    dataset = RemoraDataset.load_from_file(
        args.remora_dataset_path,
        shuffle_on_iter=False,
        drop_last=False,
    )
    print(f"Dataset summary:\n{dataset.summary}")

def register_dataset_merge(parser):
    subparser = parser.add_parser(
        "merge",
        description="Merge Remora datasets",
        help="Merge Remora datasets",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "--input-dataset",
        nargs=2,
        metavar=("PATH", "NUM_CHUNKS"),
        action="append",
        help="""1) Remora training dataset path and 2) max number of chunks
        to extract from this dataset. Second argument can be "all" to use all
        chunks from a dataset""",
    )
    subparser.add_argument(
        "--output-dataset",
        required=True,
        help="Output path for dataset",
    )
    subparser.add_argument(
        "--balance",
        action="store_true",
        help="Automatically balance classes after merging",
    )
    subparser.add_argument(
        "--log-filename",
        help="Log filename. (default: Don't output log file)",
    )
    subparser.set_defaults(func=run_dataset_merge)

def run_dataset_merge(args):
    from remora.data_chunks import merge_datasets

    if args.log_filename is not None:
        log.init_logger(args.log_filename)

    input_datasets = [
        (ds_path, None if num_chunks == "all" else int(num_chunks))
        for ds_path, num_chunks in args.input_dataset
    ]
    output_dataset = merge_datasets(input_datasets, args.balance)
    output_dataset.save(args.output_dataset)

################
# remora model #
################

def register_model(parser):
    subparser = parser.add_parser(
        "model",
        description="Remora model operations",
        help="Train or perform operations on Remora models",
        formatter_class=SubcommandHelpFormatter,
    )
    ssubparser = subparser.add_subparsers(title="model commands")
    #  Since `model` has several sub-commands, print help as default
    subparser.set_defaults(func=lambda x: subparser.print_help())
    #  Register model sub commands
    register_model_train(ssubparser)
    register_model_export(ssubparser)
    register_model_list_pretrained(ssubparser)
    register_model_download(ssubparser)

def register_model_train(parser):
    subparser = parser.add_parser(
        "train",
        description="Train Remora model",
        help="Train Remora model",
        formatter_class=SubcommandHelpFormatter,
    )

    subparser.add_argument(
        "remora_dataset_path",
        help="Remora training dataset",
    )

    data_grp = subparser.add_argument_group("Data Arguments")
    data_grp.add_argument(
        "--val-prop",
        default=constants.DEFAULT_VAL_PROP,
        type=float,
        help="Proportion of the dataset to be used as validation.",
    )
    data_grp.add_argument(
        "--batch-size",
        default=constants.DEFAULT_BATCH_SIZE,
        type=int,
        help="Number of samples per batch.",
    )
    data_grp.add_argument(
        "--chunk-context",
        type=int,
        nargs=2,
        metavar=("NUM_BEFORE", "NUM_AFTER"),
        help="""Override chunk context from data prep. Number of context signal
        points to select around the central position.""",
    )
    data_grp.add_argument(
        "--kmer-context-bases",
        nargs=2,
        type=int,
        metavar=("BASES_BEFORE", "BASES_AFTER"),
        help="""Override kmer context bases from data prep. Definition of
        k-mer (derived from the reference) passed into the model along with
        each signal position.""",
    )
    data_grp.add_argument(
        "--ext-val",
        nargs="+",
        help="Path(s) to the external validation Remora datasets.",
    )
    data_grp.add_argument(
        "--ext-val-names",
        nargs="+",
        help="""Names for external datasets. If provided must match length of
        [--ext-val] argument""",
    )
    data_grp.add_argument(
        "--balance",
        action="store_true",
        help="Balance classes exactly prior to training",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument(
        "--output-path",
        default="remora_train_results",
        help="Path to the output files.",
    )
    out_grp.add_argument(
        "--save-freq",
        default=10,
        type=int,
        help="After how many epochs to save the model.",
    )
    out_grp.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing output directory if existing.",
    )

    mdl_grp = subparser.add_argument_group("Model Arguments")
    mdl_grp.add_argument(
        "--model", required=True, help="Model architecture file (required)"
    )
    mdl_grp.add_argument(
        "--size",
        type=int,
        default=constants.DEFAULT_NN_SIZE,
        help="Model layer size.",
    )

    train_grp = subparser.add_argument_group("Training Arguments")
    train_grp.add_argument(
        "--epochs",
        default=constants.DEFAULT_EPOCHS,
        type=int,
        help="Number of training epochs.",
    )
    train_grp.add_argument(
        "--optimizer",
        default=constants.OPTIMIZERS[0],
        choices=constants.OPTIMIZERS,
        help="Optimizer setting.",
    )
    train_grp.add_argument(
        "--scheduler",
        default=None,
        help="Scheduler setting.",
    )
    train_grp.add_argument(
        "--lr",
        default=constants.DEFAULT_LR,
        type=float,
        help="Learning rate setting.",
    )
    train_grp.add_argument(
        "--weight-decay",
        default=constants.DEFAULT_WEIGHT_DECAY,
        type=float,
        help="Weight decay setting.",
    )
    train_grp.add_argument(
        "--early-stopping",
        default=5,
        type=int,
        help="""Stops training after a number of epochs without improvement.
        If set to 0 no stopping is done.""",
    )
    train_grp.add_argument(
        "--seed",
        type=int,
        help="Seed value.",
    )
    train_grp.add_argument(
        "--filter-fraction",
        default=constants.DEFAULT_FILT_FRAC,
        type=float,
        help="""Fraction of predictions to filter in validation reporting.
        Un-filtered validation metrics will always be reported as well.""",
    )
    train_grp.add_argument(
        "--lr-sched-kwargs",
        nargs=3,
        action="append",
        metavar=("NAME", "VALUE", "TYPE"),
    )
    train_grp.add_argument(
        "--batch-label-weights",
        type=float,
        nargs="+",
        help="Select batch labels with specified weights",
    )
    train_grp.add_argument(
        "--high-conf-incorrect-thr-frac",
        nargs=2,
        type=float,
        metavar=("THRESHOLD", "FRACTION"),
        help="1.) Threshold value of what to consider a high confidence "
        " predicition. Based on the softmax output. \n"
        "2.) Fraction (of the batch size) of highly confident incorrect "
        "predictions to filter during training. Filtering up to this value, "
        "but might be lower.",
    )
    train_grp.add_argument(
        "--finetune-path",
        help="Path to the torch checkpoint for the model to be fine tuned.",
    )
    train_grp.add_argument(
        "--freeze-num-layers",
        default=0,
        type=int,
        help="Number of layers to be frozen for finetuning.",
    )

    comp_grp = subparser.add_argument_group("Compute Arguments")
    comp_grp.add_argument(
        "--device",
        help="Device for neural network processing. See torch.device.",
    )

    subparser.set_defaults(func=run_model_train)

def run_model_train(args):
    from remora.util import parse_device
    from remora.train_model import train_model

    out_path = Path(args.output_path)
    if args.overwrite:
        if out_path.is_dir():
            rmtree(out_path)
        elif out_path.exists():
            out_path.unlink()
    elif out_path.exists():
        raise RemoraError("Refusing to overwrite existing training directory.")
    out_path.mkdir(parents=True, exist_ok=True)
    log.init_logger(os.path.join(out_path, "log.txt"))
    train_model(
        args.seed,
        parse_device(args.device),
        out_path,
        args.remora_dataset_path,
        args.chunk_context,
        args.kmer_context_bases,
        args.val_prop,
        args.batch_size,
        args.model,
        args.size,
        args.optimizer,
        args.lr,
        args.scheduler,
        args.weight_decay,
        args.epochs,
        args.save_freq,
        args.early_stopping,
        args.filter_fraction,
        args.ext_val,
        args.ext_val_names,
        args.lr_sched_kwargs,
        args.balance,
        args.batch_label_weights,
        args.high_conf_incorrect_thr_frac,
        args.finetune_path,
        args.freeze_num_layers,
    )

def register_model_export(parser):
    subparser = parser.add_parser(
        "export",
        description="Export a model to TorchScript format for inference.",
        help="Export a model to TorchScript format for inference.",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "checkpoint_path",
        help="Path to a pretrained model checkpoint.",
    )
    subparser.add_argument(
        "output_path",
        help="Path or directory to save the model.",
    )
    subparser.add_argument(
        "--model-path",
        help="Path to a model architecture. Default: Use path from checkpoint.",
    )
    subparser.add_argument(
        "--format",
        default="dorado",
        choices=["dorado", "torchscript"],
        help="Export format.",
    )

    subparser.set_defaults(func=run_model_export)

def run_model_export(args):
    from remora.model_util import (
        continue_from_checkpoint,
        load_torchscript_model,
        export_model_dorado,
        export_model_torchscript,
    )

    LOGGER.info("Loading model")
    try:
        model, ckpt = load_torchscript_model(args.checkpoint_path)
        LOGGER.info("Loaded a torchscript model")
    except RuntimeError:
        ckpt, model = continue_from_checkpoint(
            args.checkpoint_path, args.model_path
        )
        LOGGER.info("Loaded model from checkpoint")

    if args.format == "dorado":
        LOGGER.info("Exporting model to dorado format")
        export_model_dorado(ckpt, model, args.output_path)
    elif args.format == "torchscript":
        LOGGER.info("Exporting model to TorchScript format")
        export_model_torchscript(ckpt, model, args.output_path)
    else:
        raise RemoraError(f"Invalid export format: {args.format}")

def register_model_list_pretrained(parser):
    subparser = parser.add_parser(
        "list_pretrained",
        description="List pre-trained modified base models.",
        help="List pre-trained modified base models.",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument("--pore", help="specify pore type")
    subparser.add_argument(
        "--basecall-model-type",
        help="Specify the basecaller model type (e.g., fast, hac or sup)",
    )
    subparser.add_argument(
        "--basecall-model-version", help="Specify basecaller model version"
    )
    subparser.add_argument(
        "--modified-bases",
        nargs="+",
        help="Specify the modified base(s)",
    )
    subparser.add_argument(
        "--remora-model-type",
        help="Specify model motif (sequence context)",
    )
    subparser.add_argument(
        "--remora-model-version", help="Specify Remora model version"
    )
    subparser.set_defaults(func=run_list_pretrained)

def run_list_pretrained(args):
    from remora.model_util import get_pretrained_models
    from tabulate import tabulate

    models, header = get_pretrained_models(
        args.pore,
        args.basecall_model_type,
        args.basecall_model_version,
        args.modified_bases,
        args.remora_model_type,
        args.remora_model_version,
    )
    LOGGER.info(
        "Remora pretrained modified base models:\n"
        + tabulate(models, headers=header, showindex=False)
    )

def register_model_download(parser):
    subparser = parser.add_parser(
        "download",
        description="Download pre-trained modified base models.",
        help="Download pre-trained modified base models.",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument("--pore", help="specify pore type")
    subparser.add_argument(
        "--basecall-model-type",
        help="Specify the basecaller model type (e.g., fast, hac or sup)",
    )
    subparser.add_argument(
        "--basecall-model-version", help="Specify basecaller model version"
    )
    subparser.add_argument(
        "--modified-bases",
        nargs="+",
        help="Specify the modified base(s)",
    )
    subparser.add_argument(
        "--remora-model-type",
        help="Specify model motif (sequence context)",
    )
    subparser.add_argument(
        "--remora-model-version", help="Specify Remora model version"
    )
    subparser.set_defaults(func=run_download)

def run_download(args):
    import pkg_resources

    from remora.download import ModelDownload
    from remora.model_util import get_pretrained_models

    models, header = get_pretrained_models(
        args.pore,
        args.basecall_model_type,
        args.basecall_model_version,
        args.modified_bases,
        args.remora_model_type,
        args.remora_model_version,
    )
    path = pkg_resources.resource_filename(
        "remora",
        constants.MODEL_DATA_DIR_NAME,
    )
    out_path = Path(path)
    out_path.mkdir(parents=True, exist_ok=True)
    model_dl = ModelDownload(out_path)
    for model_url in models["Remora_Model_URL"]:
        if model_url != "":
            model_dl.download(model_url)

################
# remora infer #
################

def register_infer(parser):
    subparser = parser.add_parser(
        "infer",
        description="Perform Remora model inference",
        help="Perform Remora model inference",
        formatter_class=SubcommandHelpFormatter,
    )
    ssubparser = subparser.add_subparsers(title="infer commands")
    # Since `infer` has several sub-commands, print help as default
    subparser.set_defaults(func=lambda x: subparser.print_help())
    # Register infer sub commands
    register_infer_from_pod5_and_bam(ssubparser)
    register_infer_duplex_from_pod5_and_bam(ssubparser)

def register_infer_from_pod5_and_bam(parser):
    subparser = parser.add_parser(
        "from_pod5_and_bam",
        description="Infer modified bases from pod5 and bam inputs",
        help="Run inference on pod5s and alignments",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "pod5",
        help="POD5 file corresponding to bam file.",
    )
    subparser.add_argument(
        "in_bam",
        help="BAM file containing mv tags.",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument("--out-bam", help="Output BAM path.", required=True)
    out_grp.add_argument(
        "--log-filename",
        help="Log filename. Default: Don't output log file.",
    )

    mdl_grp = subparser.add_argument_group("Model Arguments")
    mdl_grp.add_argument(
        "--model",
        help="Path to a pretrained model in torchscript format.",
    )
    mdl_grp.add_argument(
        "--pore",
        help="Choose the type of pore the Remora model has been trained on "
        "(e.g. dna_r10.4_e8.1)",
    )
    mdl_grp.add_argument(
        "--basecall-model-type",
        help="Choose the basecaller model type (choose from fast, hac or sup)",
    )
    mdl_grp.add_argument(
        "--basecall-model-version",
        help="Choose a specific basecaller version",
    )
    mdl_grp.add_argument(
        "--modified-bases",
        nargs="+",
        help="Long name of the modified bases to call (e.g., 5mc, 5hmc).",
    )
    mdl_grp.add_argument(
        "--remora-model-type",
        help="""Choose the specific motif of the model you want to load.
        If None, load CG model.""",
    )
    mdl_grp.add_argument(
        "--remora-model-version",
        type=int,
        help="Choose the remora model version. If None, use latest.",
    )

    data_grp = subparser.add_argument_group("Data Arguments")
    data_grp.add_argument(
        "--num-reads",
        default=None,
        type=int,
        help="Number of reads.",
    )
    data_grp.add_argument(
        "--reference-anchored",
        action="store_true",
        help="""Infer per-read modified bases against reference bases instead
        of basecalls.""",
    )

    comp_grp = subparser.add_argument_group("Compute Arguments")
    comp_grp.add_argument(
        "--device",
        type=int,
        help="ID of GPU that is used for inference. Default: CPU only",
    )
    comp_grp.add_argument(
        "--queue-max",
        type=int,
        default=1_000,
        help="Maximum number of reads to store in each multiprocessing queue.",
    )
    comp_grp.add_argument(
        "--num-extract-alignment-workers",
        type=int,
        default=1,
        help="Number of signal extraction workers.",
    )
    comp_grp.add_argument(
        "--num-prepare-batch-workers",
        type=int,
        default=1,
        help="Number of batch preparation workers.",
    )
    comp_grp.add_argument(
        "--batch-size",
        default=constants.DEFAULT_BATCH_SIZE,
        type=int,
        help="Number of input units per batch.",
    )

    subparser.set_defaults(func=run_infer_from_pod5_and_bam)

def register_infer_duplex_from_pod5_and_bam(parser):
    duplex_delim_flag = "--duplex-delim"
    subparser = parser.add_parser(
        "duplex_from_pod5_and_bam",
        description="Infer modified bases on duplex reads from pod5 and bam "
        "inputs",
        help="Run inference on pod5s simplex reads and duplex alignments with "
        "duplex pairs",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "pod5",
        help="POD5 file corresponding to bam file.",
    )
    subparser.add_argument(
        "simplex_bam",
        help="Base called BAM file containing mv tags.",
    )
    subparser.add_argument(
        "duplex_bam",
        help=f"""BAM file containing duplex base called sequences (and optional
        reference mappings). Record names may either be the template read_id
        or template<delim>complement. The value of <delim> can be set with
        {duplex_delim_flag}.""",
    )
    subparser.add_argument(
        "duplex_read_pairs",
        help="""Whitespace separated plain text file containing read ID pairs,
        no header.""",
    )
    subparser.add_argument(
        duplex_delim_flag,
        default=";",
        help="""Deliminator string between template and complement read
        ids in the duplex BAM""",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument(
        "--out-bam",
        help="Output BAM path.",
    )
    out_grp.add_argument(
        "--log-filename",
        help="Log filename. Default: Don't output log file.",
    )

    mdl_grp = subparser.add_argument_group("Model Arguments")
    mdl_grp.add_argument(
        "--model",
        help="Path to a pretrained model in torchscript format.",
    )
    mdl_grp.add_argument(
        "--pore",
        help="""Choose the type of pore the Remora model has been trained on
        (e.g. dna_r10.4_e8.1)""",
    )
    mdl_grp.add_argument(
        "--basecall-model-type",
        help="Choose the basecaller model type (choose from fast, hac or sup)",
    )
    mdl_grp.add_argument(
        "--basecall-model-version",
        help="Choose a specific basecaller version",
    )
    mdl_grp.add_argument(
        "--modified-bases",
        nargs="+",
        help="Long name of the modified bases to call (e.g., 5mc, 5hmc).",
    )
    mdl_grp.add_argument(
        "--remora-model-type",
        help="Choose the specific motif of the model you want to load. "
        "If None, load CG model.",
    )
    mdl_grp.add_argument(
        "--remora-model-version",
        type=int,
        help="Choose the remora model version. If None, use latest.",
    )

    data_grp = subparser.add_argument_group("Data Arguments")
    data_grp.add_argument(
        "--num-reads",
        type=int,
        help="Number of reads.",
    )

    comp_grp = subparser.add_argument_group("Compute Arguments")
    comp_grp.add_argument(
        "--device",
        type=int,
        help="ID of GPU that is used for inference. Default: CPU only",
    )
    comp_grp.add_argument(
        "--num-extract-alignment-workers",
        type=int,
        default=1,
        help="Number of IO extraction workers.",
    )
    comp_grp.add_argument(
        "--num-duplex-prep-workers",
        type=int,
        default=1,
        help="Number of duplex prep workers (tends to bottleneck).",
    )
    comp_grp.add_argument(
        "--num-infer-workers",
        type=int,
        default=1,
        help="""Number of chunk extraction workers. If performing signal
        refinement this should be increased.""",
    )

    subparser.set_defaults(func=run_infer_from_pod5_and_bam_duplex)

def _unpack_model_kw_args(args) -> dict:
    from remora.util import parse_device

    if args.model and not os.path.exists(args.model):
        raise ValueError(f"didn't find model file at {args.model}")

    model_kwargs = {
        "model_filename": args.model,
        "pore": args.pore,
        "basecall_model_type": args.basecall_model_type,
        "basecall_model_version": args.basecall_model_version,
        "modified_bases": args.modified_bases,
        "remora_model_type": args.remora_model_type,
        "remora_model_version": args.remora_model_version,
        "device": parse_device(args.device),
    }
    return model_kwargs

def run_infer_from_pod5_and_bam(args):
    import torch

    from remora.model_util import load_model
    from remora.inference import infer_from_pod5_and_bam

    if args.log_filename is not None:
        log.init_logger(args.log_filename)
    # test that model can be loaded in parent process
    model_kwargs = _unpack_model_kw_args(args)
    model, model_metadata = load_model(
        **model_kwargs, quiet=False, eval_only=True
    )
    torch.set_grad_enabled(False)
    infer_from_pod5_and_bam(
        pod5_path=args.pod5,
        in_bam_path=args.in_bam,
        model=model,
        model_metadata=model_metadata,
        out_bam_path=args.out_bam,
        num_reads=args.num_reads,
        queue_max=args.queue_max,
        num_extract_alignment_workers=args.num_extract_alignment_workers,
        num_prep_batch_workers=args.num_prepare_batch_workers,
        batch_size=args.batch_size,
        ref_anchored=args.reference_anchored,
    )

def run_infer_from_pod5_and_bam_duplex(args):
    import torch

    from remora.model_util import load_model
    from remora.inference import infer_duplex

    if args.log_filename is not None:
        log.init_logger(args.log_filename)
    model_kwargs = _unpack_model_kw_args(args)
    model, model_metadata = load_model(
        **model_kwargs, quiet=False, eval_only=True
    )
    torch.set_grad_enabled(False)

    if not os.path.exists(args.pod5):
        raise ValueError(f"didn't find pod5 at {args.pod5}")
    if not os.path.exists(args.simplex_bam):
        raise ValueError(f"didn't find simplex bam at {args.simplex_bam}")
    if not os.path.exists(args.duplex_bam):
        raise ValueError(f"didn't find duplex bam at {args.duplex_bam}")
    if not os.path.exists(args.duplex_read_pairs):
        raise ValueError(
            f"didn't find duplex read pairs at {args.duplex_read_pairs}"
        )

    infer_duplex(
        simplex_pod5_path=args.pod5,
        simplex_bam_path=args.simplex_bam,
        duplex_bam_path=args.duplex_bam,
        pairs_path=args.duplex_read_pairs,
        model=model,
        model_metadata=model_metadata,
        out_bam=args.out_bam,
        num_extract_alignment_threads=args.num_extract_alignment_workers,
        num_duplex_prep_workers=args.num_duplex_prep_workers,
        num_infer_threads=args.num_infer_workers,
        num_reads=args.num_reads,
        duplex_deliminator=args.duplex_delim,
    )

###################
# remora validate #
###################

def register_validate(parser):
    subparser = parser.add_parser(
        "validate",
        description="Validate modified base predictions",
        help="Validate modified base predictions",
        formatter_class=SubcommandHelpFormatter,
    )
    ssubparser = subparser.add_subparsers(title="validation commands")
    # Since `validate` has several sub-commands, print help as default
    subparser.set_defaults(func=lambda x: subparser.print_help())
    # Register validate sub commands
    register_validate_from_modbams(ssubparser)
    register_validate_from_remora_dataset(ssubparser)

def register_validate_from_modbams(parser):
    subparser = parser.add_parser(
        "from_modbams",
        description="Validation with ground truth samples",
        help="Validation with ground truth samples",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "--bam-and-bed",
        required=True,
        nargs=2,
        metavar=("BAM", "GROUND_TRUTH_BED"),
        action="append",
        help="""Argument accepts 2 values. The first value is the BAM file path
        with modified base tags. The second is a bed file with ground truth
        reference positions. The name field in the ground truth bed file should
        be the single letter code for a modified base or the corresponding
        canonical base. This argument can be provided more than once for
        multiple samples.""",
    )
    subparser.add_argument(
        "--full-results-filename", help="Output per-read calls to TSV file."
    )
    subparser.add_argument(
        "--name",
        default="sample",
        help="""Name of this sample/comparison. Useful when tabulating
        several runs.""",
    )
    subparser.add_argument(
        "--pct-filt",
        type=float,
        default=10.0,
        help="Filter a specified percentage (or less given ties) of calls.",
    )
    subparser.add_argument(
        "--allow-unbalanced",
        action="store_true",
        help="Allow classes to be unbalanced for metric computation.",
    )
    subparser.add_argument(
        "--max-sites-per-read",
        type=int,
        help="Maxiumum number of sites to extract from a single read.",
    )
    subparser.add_argument(
        "--seed",
        type=int,
        help="Seed value. Default: Random seed",
    )
    subparser.add_argument(
        "--extra-bases",
        help="Extra canoncial or modified base single letter codes not in "
        "the ground truth bed files which should be added to the accepted "
        "alphabet. For example, to run a sample with canonical ground truth "
        "(C) and 5mC and 5hmC calls (m and h) modified base calls this "
        "argument would be `--extra-bases mh`",
    )
    subparser.add_argument(
        "--log-filename",
        help="Log filename. (default: Don't output log file)",
    )

    subparser.set_defaults(func=run_validate_modbams)

def run_validate_modbams(args):
    from remora.validate import validate_modbams

    if args.log_filename is not None:
        log.init_logger(args.log_filename)
    validate_modbams(
        bams_and_beds=args.bam_and_bed,
        full_results_path=args.full_results_filename,
        name=args.name,
        pct_filt=args.pct_filt,
        allow_unbalanced=args.allow_unbalanced,
        seed=args.seed,
        extra_bases=args.extra_bases,
        max_sites_per_read=args.max_sites_per_read,
    )

def register_validate_from_remora_dataset(parser):
    subparser = parser.add_parser(
        "from_remora_dataset",
        description="Run validation on external Remora dataset",
        help="Validate on Remora dataset",
        formatter_class=SubcommandHelpFormatter,
    )
    subparser.add_argument(
        "remora_dataset_path",
        help="Remora training dataset",
    )

    mdl_grp = subparser.add_argument_group("Model Arguments")
    mdl_grp.add_argument(
        "--model",
        help="Path to a pretrained model.",
    )
    mdl_grp.add_argument(
        "--pore",
        help="Choose the type of pore the Remora model has been trained on "
        "(e.g. dna_r10.4_e8.1)",
    )
    mdl_grp.add_argument(
        "--basecall-model-type",
        help="Choose the basecaller model type (choose from fast, hac or sup)",
    )
    mdl_grp.add_argument(
        "--basecall-model-version",
        help="Choose a specific basecaller version",
    )
    mdl_grp.add_argument(
        "--modified-bases",
        nargs="+",
        help="Long name of the modified bases to call (e.g., 5mc, 5hmc).",
    )
    mdl_grp.add_argument(
        "--remora-model-type",
        help="Choose the specific motif of the model you want to load. "
        "If None, load CG model.",
    )
    mdl_grp.add_argument(
        "--remora-model-version",
        type=int,
        help="Choose the remora model version. If None, use latest.",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument(
        "--out-file",
        help="Output path for the validation result file.",
    )
    out_grp.add_argument(
        "--full-results-filename", help="Output per-read calls to TSV file."
    )

    val_grp = subparser.add_argument_group("Validation Arguments")
    val_grp.add_argument(
        "--pct-filt",
        type=float,
        default=10.0,
        help="Filter a specified percentage (or less given ties) of calls.",
    )

    comp_grp = subparser.add_argument_group("Compute Arguments")
    comp_grp.add_argument(
        "--batch-size",
        default=constants.DEFAULT_BATCH_SIZE,
        type=int,
        help="Number of input units per batch.",
    )
    comp_grp.add_argument(
        "--device",
        type=int,
        help="ID of GPU that is used for inference. Default: CPU",
    )

    subparser.set_defaults(func=run_validate_from_remora_dataset)

def run_validate_from_remora_dataset(args):
    import torch

    from remora.util import parse_device
    from remora.model_util import load_model
    from remora.validate import ValidationLogger
    from remora.data_chunks import RemoraDataset

    LOGGER.info("Loading dataset from Remora file")
    dataset = RemoraDataset.load_from_file(
        args.remora_dataset_path,
        batch_size=args.batch_size,
        shuffle_on_iter=False,
        drop_last=False,
    )

    LOGGER.info("Loading model")
    model, model_metadata = load_model(
        args.model,
        pore=args.pore,
        basecall_model_type=args.basecall_model_type,
        basecall_model_version=args.basecall_model_version,
        modified_bases=args.modified_bases,
        remora_model_type=args.remora_model_type,
        remora_model_version=args.remora_model_version,
        device=parse_device(args.device),
        eval_only=True,
    )
    torch.set_grad_enabled(False)

    dataset.trim_kmer_context_bases(model_metadata["kmer_context_bases"])
    dataset.trim_chunk_context(model_metadata["chunk_context"])
    LOGGER.info(f"Loaded dataset summary:\n{dataset.summary}")

    if args.out_file is None:
        out_fp = sys.stdout
    else:
        out_fp = open(args.out_file, "w", buffering=1)
        atexit.register(out_fp.close)
    if args.full_results_filename is None:
        full_results_fp = None
    else:
        full_results_fp = open(args.full_results_filename, "w", buffering=1)
        atexit.register(full_results_fp.close)

    LOGGER.info("Running validation")
    val_fp = ValidationLogger(out_fp, full_results_fp)
    val_fp.validate_model(
        model,
        model_metadata["mod_bases"],
        torch.nn.CrossEntropyLoss(),
        dataset,
        args.pct_filt / 100,
    )

##################
# remora analyze #
##################

def register_analyze(parser):
    subparser = parser.add_parser(
        "analyze",
        description="Analyze nanopore data including raw signal",
        help="Analyze nanopore data including raw signal",
        formatter_class=SubcommandHelpFormatter,
    )
    ssubparser = subparser.add_subparsers(title="Analyze commands")
    # register_estimate_kmer_levels(ssubparser)
    # Since `plot` has several sub-commands, print help as default
    subparser.set_defaults(func=lambda x: subparser.print_help())
    # Register analyze sub commands
    register_analyze_plot(ssubparser)

def register_analyze_plot(parser):
    subparser = parser.add_parser(
        "plot",
        description="Plot nanopore data including raw signal",
        help="Plot nanopore data including raw signal",
        formatter_class=SubcommandHelpFormatter,
    )
    ssubparser = subparser.add_subparsers(title="Plot commands")
    # Since `plot` has several sub-commands, print help as default
    subparser.set_defaults(func=lambda x: subparser.print_help())
    # Register plot sub commands
    register_plot_ref_region(ssubparser)

def register_plot_ref_region(parser):
    subparser = parser.add_parser(
        "ref_region",
        description="Plot signal at reference region",
        help="Plot signal at reference region",
        formatter_class=SubcommandHelpFormatter,
    )

    in_grp = subparser.add_argument_group("Input Arguments")
    in_grp.add_argument(
        "--pod5-and-bam",
        required=True,
        nargs=2,
        metavar=("POD5", "BAM"),
        action="append",
        help="""POD5 signal path and BAM file path. BAM file must be mapped,
        sorted and indexed and contain move table and MD tags. Multiple
        samples can be supplied and will be plotted in different colors""",
    )
    in_grp.add_argument(
        "--ref-regions",
        required=True,
        help="""Reference region(s) to plot specified in BED format. Each line
        in this input file will produce one page in the output PDF.""",
    )
    in_grp.add_argument(
        "--highlight-ranges",
        help="""BED file containing regions to highlight""",
    )
    in_grp.add_argument(
        "--highlight-color",
        default="orange",
        help="""Color or highlighted regions""",
    )

    refine_grp = subparser.add_argument_group("Signal Mapping Refine Arguments")
    refine_grp.add_argument(
        "--refine-kmer-level-table",
        help="""Tab-delimited file containing no header and two fields:
        1. string k-mer sequence and 2. float expected normalized level.
        All k-mers must be the same length and all combinations of the bases
        'ACGT' must be present in the file.""",
    )
    refine_grp.add_argument(
        "--refine-rough-rescale",
        action="store_true",
        help="""Apply a rough rescaling using quantiles of signal+move table
        and levels.""",
    )
    refine_grp.add_argument(
        "--refine-scale-iters",
        default=0,
        type=int,
        help="""Number of iterations of signal mapping refinement and signal
        re-scaling to perform. Set to 0 (default) in order to perform signal
        mapping refinement, but skip re-scaling. Set to -1 to skip signal
        mapping (potentially using levels for rough rescaling).""",
    )
    refine_grp.add_argument(
        "--refine-half-bandwidth",
        default=constants.DEFAULT_REFINE_HBW,
        type=int,
        help="""Half bandwidth around signal mapping over which to search for
        new path.""",
    )
    refine_grp.add_argument(
        "--refine-algo",
        default=constants.DEFAULT_REFINE_ALGO,
        choices=constants.REFINE_ALGOS,
        help="Refinement algorithm to apply (if kmer level table is provided).",
    )
    refine_grp.add_argument(
        "--refine-short-dwell-parameters",
        default=constants.DEFAULT_REFINE_SHORT_DWELL_PARAMS,
        type=float,
        nargs=3,
        metavar=("TARGET", "LIMIT", "WEIGHT"),
        help="""Short dwell penalty refiner parameters. Dwells shorter than
        LIMIT will be penalized a value of `WEIGHT * (dwell - TARGET)^2`.""",
    )

    plt_grp = subparser.add_argument_group("Plot Arguments")
    plt_grp.add_argument(
        "--figsize",
        nargs=2,
        type=int,
        metavar=("HEIGHT", "WIDTH"),
        default=constants.DEFAULT_PLOT_FIG_SIZE,
        help="Figure size",
    )
    plt_grp.add_argument(
        "--ylim",
        nargs=2,
        type=int,
        metavar=("MIN", "MAX"),
        help="Signal plotting limits",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument(
        "--plots-filename",
        default="remora_raw_signal_plot.pdf",
        help="Output plots PDF file location.",
    )
    out_grp.add_argument(
        "--log-filename",
        help="Log filename. Default: Don't output log file.",
    )

    subparser.set_defaults(func=run_plot_ref_region)

def run_plot_ref_region(args):
    import pod5
    import pysam
    from matplotlib import pyplot as plt
    from matplotlib.backends.backend_pdf import PdfPages

    from remora import log, io, refine_signal_map

    if args.log_filename is not None:
        log.init_logger(args.log_filename)
    pod5_paths, bc_paths = zip(*args.pod5_and_bam)
    bam_fhs = [pysam.AlignmentFile(bc_path) for bc_path in bc_paths]
    pod5_fhs = [pod5.Reader(pod5_path) for pod5_path in pod5_paths]
    sig_map_refiner = refine_signal_map.SigMapRefiner(
        kmer_model_filename=args.refine_kmer_level_table,
        do_rough_rescale=args.refine_rough_rescale,
        scale_iters=args.refine_scale_iters,
        algo=args.refine_algo,
        half_bandwidth=args.refine_half_bandwidth,
        sd_params=args.refine_short_dwell_parameters,
        do_fix_guage=True,
    )
    highlight_ranges = None
    if args.highlight_ranges is not None:
        highlight_ranges = io.parse_bed(args.highlight_ranges)

    with PdfPages(args.plots_filename) as pdf_fh:
        for ref_reg in io.parse_bed_lines(args.ref_regions):
            reg_highlight_ranges = None
            if highlight_ranges is not None:
                try:
                    reg_highlight_ranges = [
                        (pos, pos + 1, args.highlight_color)
                        for pos in highlight_ranges[
                            (ref_reg.ctg, ref_reg.strand)
                        ]
                        if ref_reg.start <= pos < ref_reg.end
                    ]
                except KeyError:
                    LOGGER.debug(f"No highlight values for region {ref_reg}")
                    pass
            fig, ax = plt.subplots(figsize=args.figsize)
            io.plot_signal_at_ref_region(
                ref_reg,
                zip(pod5_fhs, bam_fhs),
                sig_map_refiner,
                fig_ax=(fig, ax),
                ylim=args.ylim,
                highlight_ranges=reg_highlight_ranges,
            )
            pdf_fh.savefig(fig, bbox_inches="tight")

def register_estimate_kmer_levels(parser):
    subparser = parser.add_parser(
        "estimate_kmer_levels",
        description="Estimate k-mer level table",
        help="Estimate k-mer level table",
        formatter_class=SubcommandHelpFormatter,
    )

    in_grp = subparser.add_argument_group("Input Arguments")
    in_grp.add_argument(
        "--pod5-and-bam",
        required=True,
        nargs=2,
        metavar=("POD5", "BAM"),
        action="append",
        help="""POD5 signal path and BAM file path. BAM file must be mapped,
        sorted and indexed and contain move table and MD tags. Multiple
        samples can be supplied and will be aggregated after site level
        extraction""",
    )

    refine_grp = subparser.add_argument_group("Signal Mapping Refine Arguments")
    refine_grp.add_argument(
        "--refine-kmer-level-table",
        help="""Tab-delimited file containing no header and two fields:
        1. string k-mer sequence and 2. float expected normalized level.
        All k-mers must be the same length and all combinations of the bases
        'ACGT' must be present in the file.""",
    )
    refine_grp.add_argument(
        "--refine-rough-rescale",
        action="store_true",
        help="""Apply a rough rescaling using quantiles of signal+move table
        and levels.""",
    )
    refine_grp.add_argument(
        "--refine-scale-iters",
        default=0,
        type=int,
        help="""Number of iterations of signal mapping refinement and signal
        re-scaling to perform. Set to 0 (default) in order to perform signal
        mapping refinement (aka resquiggle), but skip fine re-scaling. Set to
        -1 to skip signal mapping (potentially using levels for rough
        rescaling).""",
    )
    refine_grp.add_argument(
        "--refine-half-bandwidth",
        default=constants.DEFAULT_REFINE_HBW,
        type=int,
        help="""Half bandwidth around signal mapping over which to search for
        new path.""",
    )
    refine_grp.add_argument(
        "--refine-algo",
        default=constants.DEFAULT_REFINE_ALGO,
        choices=constants.REFINE_ALGOS,
        help="Refinement algorithm to apply (if kmer level table is provided).",
    )
    refine_grp.add_argument(
        "--refine-short-dwell-parameters",
        default=constants.DEFAULT_REFINE_SHORT_DWELL_PARAMS,
        type=float,
        nargs=3,
        metavar=("TARGET", "LIMIT", "WEIGHT"),
        help="""Short dwell penalty refiner parameters. Dwells shorter than
        LIMIT will be penalized a value of `WEIGHT * (dwell - TARGET)^2`.""",
    )

    data_grp = subparser.add_argument_group("Data Arguments")
    data_grp.add_argument(
        "--min-coverage",
        type=int,
        default=10,
        help="Miniumum coverage to include a site.",
    )
    data_grp.add_argument(
        "--kmer-context-bases",
        nargs=2,
        default=(2, 2),
        type=int,
        metavar=("BASES_BEFORE", "BASES_AFTER"),
        help="""Definition of k-mer by the number of bases before and after the
        assigned signal position""",
    )

    out_grp = subparser.add_argument_group("Output Arguments")
    out_grp.add_argument(
        "--levels-filename",
        default="remora_kmer_levels.txt",
        help="Output file for kmer levels.",
    )
    out_grp.add_argument(
        "--log-filename",
        help="Log filename. Default: Don't output log file.",
    )

    comp_grp = subparser.add_argument_group("Compute Arguments")
    comp_grp.add_argument(
        "--num-workers",
        type=int,
        default=1,
        help="Number of workers.",
    )
    comp_grp.add_argument(
        "--chunk-width",
        type=int,
        default=1_000,
        help="""Width of reference region to process at one time. Should be
        smaller for very high coverage.""",
    )
    comp_grp.add_argument(
        "--max-chunk-coverage",
        type=int,
        default=100,
        help="Maxiumum mean chunk coverage for each region.",
    )

    subparser.set_defaults(func=run_estimate_kmer_levels)

def run_estimate_kmer_levels(args):
    from itertools import product

    import pysam
    import numpy as np

    from remora import log, io, refine_signal_map

    if args.log_filename is not None:
        log.init_logger(args.log_filename)
    # open first to avoid long process without write access
    out_fh = open(args.levels_filename, "w")

    sig_map_refiner = refine_signal_map.SigMapRefiner(
        kmer_model_filename=args.refine_kmer_level_table,
        do_rough_rescale=args.refine_rough_rescale,
        scale_iters=args.refine_scale_iters,
        algo=args.refine_algo,
        half_bandwidth=args.refine_half_bandwidth,
        sd_params=args.refine_short_dwell_parameters,
        do_fix_guage=True,
    )
    if not sig_map_refiner.is_loaded or sig_map_refiner.scale_iters < 0:
        LOGGER.warning(
            "It is highly recommended to apply signal mapping refinement in "
            "order to output a valid kmer level table."
        )

    kmer_len = sum(args.kmer_context_bases) + 1
    all_kmer_levels = dict(
        ("".join(bs), []) for bs in product("ACGT", repeat=kmer_len)
    )
    for pod5_path, bam_path in args.pod5_and_bam:
        try:
            with pysam.AlignmentFile(bam_path) as bam_fh:
                _ = bam_fh.fetch(bam_fh.header.references[0], 0, 1)
        except ValueError:
            LOGGER.warning(
                "Cannot estimate levels from BAM file without mappings or index"
            )
            continue
        LOGGER.info(f"Extracting levels from {pod5_path} and {bam_path}")
        for kmer, levels in io.get_site_kmer_levels(
            pod5_path,
            bam_path,
            sig_map_refiner,
            args.kmer_context_bases,
            min_cov=args.min_coverage,
            chunk_len=args.chunk_width,
            max_chunk_cov=args.max_chunk_coverage,
            num_workers=args.num_workers,
        ).items():
            all_kmer_levels[kmer].append(levels)
    LOGGER.info("Aggregating and outputting levels")
    been_warned = False
    for kmer, levels in sorted(all_kmer_levels.items()):
        levels = np.concatenate(levels)
        if levels.size == 0:
            if not been_warned:
                LOGGER.warning("Some k-mers not observed.")
                been_warned = True
            out_fh.write(f"{kmer}\tnan\n")
        else:
            out_fh.write(f"{kmer}\t{np.median(levels)}\n")
    out_fh.close()


import numpy as np
import os
import pandas as pd
import glob

from torchtext.data.metrics import bleu_score
from jiwer import wer

def get_csv(root):
    root = os.path.join(root, 'phoenix2014-release/phoenix-2014-multisigner')
    root = os.path.join(root, 'annotations/manual/train.corpus.csv')
    return pd.read_csv(root)

def DatasetStatistic(root, mode='train'):
    '''
    analysis the statistic of the RWTH-PHOENIX-Weather 2014 dataset
    Args:
        root: the whole data root
        mode: 'train','dev', 'test'
    return:
        max_len: maximum video length
    '''

    root = os.path.join(root, 'phoenix2014-release/phoenix-2014-multisigner')
    video_root = os.path.join(root, 'features/fullFrame-210x260px/' + mode)
    csv_path = os.path.join(root, 'annotations/manual/' + mode + '.corpus.csv')

    csv_file = pd.read_csv(csv_path)
    video_paths = [os.path.join(video_root, csv_file.iloc[i, 0].split('|')[1])
                   for i in range(csv_file.shape[0])]

    video_lens = [len(glob.glob(path)) for path in video_paths]
    return [max(video_lens), min(video_lens),
            np.mean(video_lens), np.std(video_lens)]

def itos(idx_seq, TRG):
    '''
    Description:
        denumericalize: convert the index sequence to the text sequence
    '''
    return [TRG.vocab.itos[idx] for idx in idx_seq]

def bleu_count(outputs, targets, TRG):
    '''
    sentence level
    shape:
        outputs: [T, N, E]
        targets: [T, N]
    '''
    outputs = outputs.max(2)[1].transpose(0,1)
    targets = targets.transpose(0,1)
    candidate_corpus = [itos(idx_seq, TRG) for idx_seq in outputs]
    references_corpus = [[itos(idx_seq, TRG)] for idx_seq in targets]
    return bleu_score(candidate_corpus, references_corpus)

def wer_count(outputs, targets, TRG):
    '''
    word error rate
    Ref: https://pypi.org/project/jiwer/
    shape:
        outputs: [T, N, E]
        targets: [T, N]
    '''
    outputs = outputs.max(2)[1].transpose(0,1)
    targets = targets.transpose(0,1)
    candidate_corpus = [' '.join(itos(idx_seq, TRG)) for idx_seq in outputs]
    reference_corpus = [' '.join(itos(idx_seq, TRG)) for idx_seq in targets]
    return wer(reference_corpus, candidate_corpus, standardize = True)


#!/usr/bin/env python3
"""
@Filename:    utils.py
@Author:      dulanj
@Time:        2021-09-23 00.27
"""
import collections

import numpy as np
from matplotlib import gridspec
from matplotlib import pyplot as plt

DatasetInfo = collections.namedtuple(
    'DatasetInfo',
    'num_classes, label_divisor, thing_list, colormap, class_names')

def _cityscapes_label_colormap():
    """Creates a label colormap used in CITYSCAPES segmentation benchmark.

    See more about CITYSCAPES dataset at https://www.cityscapes-dataset.com/
    M. Cordts, et al. "The Cityscapes Dataset for Semantic Urban Scene Understanding." CVPR. 2016.

    Returns:
      A 2-D numpy array with each row being mapped RGB color (in uint8 range).
    """
    colormap = np.zeros((256, 3), dtype=np.uint8)
    colormap[0] = [128, 64, 128]
    colormap[1] = [244, 35, 232]
    colormap[2] = [70, 70, 70]
    colormap[3] = [102, 102, 156]
    colormap[4] = [190, 153, 153]
    colormap[5] = [153, 153, 153]
    colormap[6] = [250, 170, 30]
    colormap[7] = [220, 220, 0]
    colormap[8] = [107, 142, 35]
    colormap[9] = [152, 251, 152]
    colormap[10] = [70, 130, 180]
    colormap[11] = [220, 20, 60]
    colormap[12] = [255, 0, 0]
    colormap[13] = [0, 0, 142]
    colormap[14] = [0, 0, 70]
    colormap[15] = [0, 60, 100]
    colormap[16] = [0, 80, 100]
    colormap[17] = [0, 0, 230]
    colormap[18] = [119, 11, 32]
    return colormap

def _cityscapes_class_names():
    return ('road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
            'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
            'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle',
            'bicycle')

def cityscapes_dataset_information():
    return DatasetInfo(
        num_classes=19,
        label_divisor=1000,
        thing_list=tuple(range(11, 19)),
        colormap=_cityscapes_label_colormap(),
        class_names=_cityscapes_class_names())

def perturb_color(color, noise, used_colors, max_trials=50, random_state=None):
    """Pertrubs the color with some noise.

    If `used_colors` is not None, we will return the color that has
    not appeared before in it.

    Args:
      color: A numpy array with three elements [R, G, B].
      noise: Integer, specifying the amount of perturbing noise (in uint8 range).
      used_colors: A set, used to keep track of used colors.
      max_trials: An integer, maximum trials to generate random color.
      random_state: An optional np.random.RandomState. If passed, will be used to
        generate random numbers.

    Returns:
      A perturbed color that has not appeared in used_colors.
    """
    if random_state is None:
        random_state = np.random

    for _ in range(max_trials):
        random_color = color + random_state.randint(
            low=-noise, high=noise + 1, size=3)
        random_color = np.clip(random_color, 0, 255)

        if tuple(random_color) not in used_colors:
            used_colors.add(tuple(random_color))
            return random_color

    print('Max trial reached and duplicate color will be used. Please consider '
          'increase noise in `perturb_color()`.')
    return random_color

def color_panoptic_map(panoptic_prediction, dataset_info, perturb_noise):
    """Helper method to colorize output panoptic map.

    Args:
      panoptic_prediction: A 2D numpy array, panoptic prediction from deeplab
        model.
      dataset_info: A DatasetInfo object, dataset associated to the model.
      perturb_noise: Integer, the amount of noise (in uint8 range) added to each
        instance of the same semantic class.

    Returns:
      colored_panoptic_map: A 3D numpy array with last dimension of 3, colored
        panoptic prediction map.
      used_colors: A dictionary mapping semantic_ids to a set of colors used
        in `colored_panoptic_map`.
    """
    if panoptic_prediction.ndim != 2:
        raise ValueError('Expect 2-D panoptic prediction. Got {}'.format(
            panoptic_prediction.shape))

    semantic_map = panoptic_prediction // dataset_info.label_divisor
    instance_map = panoptic_prediction % dataset_info.label_divisor
    height, width = panoptic_prediction.shape
    colored_panoptic_map = np.zeros((height, width, 3), dtype=np.uint8)

    used_colors = collections.defaultdict(set)
    # Use a fixed seed to reproduce the same visualization.
    random_state = np.random.RandomState(0)

    unique_semantic_ids = np.unique(semantic_map)
    for semantic_id in unique_semantic_ids:
        semantic_mask = semantic_map == semantic_id
        if semantic_id in dataset_info.thing_list:
            # For `thing` class, we will add a small amount of random noise to its
            # correspondingly predefined semantic segmentation colormap.
            unique_instance_ids = np.unique(instance_map[semantic_mask])
            for instance_id in unique_instance_ids:
                instance_mask = np.logical_and(semantic_mask,
                                               instance_map == instance_id)
                random_color = perturb_color(
                    dataset_info.colormap[semantic_id],
                    perturb_noise,
                    used_colors[semantic_id],
                    random_state=random_state)
                colored_panoptic_map[instance_mask] = random_color
        else:
            # For `stuff` class, we use the defined semantic color.
            colored_panoptic_map[semantic_mask] = dataset_info.colormap[semantic_id]
            used_colors[semantic_id].add(tuple(dataset_info.colormap[semantic_id]))
    return colored_panoptic_map, used_colors

def vis_segmentation(image,
                     panoptic_prediction,
                     dataset_info,
                     perturb_noise=60,
                     output_name="show"):
    """Visualizes input image, segmentation map and overlay view."""
    plt.figure(figsize=(30, 20))
    grid_spec = gridspec.GridSpec(2, 2)

    ax = plt.subplot(grid_spec[0])
    plt.imshow(image)
    plt.axis('off')
    ax.set_title('input image', fontsize=20)

    ax = plt.subplot(grid_spec[1])
    panoptic_map, used_colors = color_panoptic_map(panoptic_prediction,
                                                   dataset_info, perturb_noise)
    plt.imshow(panoptic_map)
    plt.axis('off')
    ax.set_title('panoptic map', fontsize=20)

    ax = plt.subplot(grid_spec[2])
    plt.imshow(image)
    plt.imshow(panoptic_map, alpha=0.7)
    plt.axis('off')
    ax.set_title('panoptic overlay', fontsize=20)

    ax = plt.subplot(grid_spec[3])
    max_num_instances = max(len(color) for color in used_colors.values())
    # RGBA image as legend.
    legend = np.zeros((len(used_colors), max_num_instances, 4), dtype=np.uint8)
    class_names = []
    for i, semantic_id in enumerate(sorted(used_colors)):
        legend[i, :len(used_colors[semantic_id]), :3] = np.array(
            list(used_colors[semantic_id]))
        legend[i, :len(used_colors[semantic_id]), 3] = 255
        if semantic_id < dataset_info.num_classes:
            class_names.append(dataset_info.class_names[semantic_id])
        else:
            class_names.append('ignore')

    plt.imshow(legend, interpolation='nearest')
    ax.yaxis.tick_left()
    plt.yticks(range(len(legend)), class_names, fontsize=15)
    plt.xticks([], [])
    ax.tick_params(width=0.0, grid_linewidth=0.0)
    plt.grid('off')
    if output_name == "show":
        plt.show()
    else:
        plt.savefig(f'{output_name}.png')


from webob import Request, Response
import datetime
from MappingCommon import MappingCommon

def application(environ, start_response):
    req = Request(environ)
    res = Response()
    res.content_type = 'text/html'

    now = str(datetime.datetime.today())

    mapc = MappingCommon()
    logFilePath = mapc.projectRoot + "/log"
    kmlMapHeight = int(mapc.getConfiguration('KMLMapHeight'))
    mapUrl = mapc.getConfiguration('MapUrl')
    instructions = mapc.getConfiguration('KMLInstructions')
    snapTolerance = mapc.getConfiguration('SnapTolerance')
    select = mapc.buildSelect()
    # initiate refJson and workJson as empty
    refJson = '""'
    workJson = '""'

    k = open(logFilePath + "/OL.log", "a")
    k.write("\ngetkml: datetime = %s\n" % now)

    kmlName = req.params['kmlName']

    try:
        dlon = float(req.params['dlon'])
        dlat = float(req.params['dlat'])
    except:
        dlon = float(mapc.getConfiguration('KMLdlon'))
        dlat = float(mapc.getConfiguration('KMLdlat'))

    gridJson = mapc.getGridJson(kmlName, dlon, dlat)
    gridJson2 = mapc.getGridJson(kmlName, dlon*2.56, dlat*2.56)

    if len(kmlName) > 0:
        (kmlType, kmlTypeDescr) = mapc.getKmlType(kmlName)
        mapHint = mapc.querySingleValue(
            "select hint from kml_data where name = '%s'" % kmlName
        )
        instanceid = mapc.querySingleValue(
            "select value from configuration where key='instance%s'" % kmlType
        )
        image_attributes = mapc.get_image_attributes(kmlName)
        # Training and field mapping cases.
        # These have an assignmentId.
        try:
            assignmentId = req.params['assignmentId']
            resultsAccepted = req.params['resultsAccepted']
            submitTo = req.params['submitTo']
            csrfToken = req.params['csrfToken']
            workerId = ''
            target = '_parent'
            commentsVisible = 'block'

            # Training case.
            # This has a tryNum.
            try:
                tryNum = req.params['tryNum']
                hitId = ''
                commentsDisabled = 'disabled'
                mapHint = '
Hint: %s
' % mapHint
                kmlMapHeight -= 30        # Reduce map height to leave room for hints.

            # Field mapping case.
            except:
                tryNum = ''
                hitId = req.params['hitId']
                commentsDisabled = ''
                mapHint = ''

        # Worker feedback and standalone cases.
        # These have no assignmentId.
        except:
            assignmentId = ''
            tryNum = ''
            hitId = ''
            resultsAccepted = ''
            submitTo = ''
            csrfToken = ''
            commentsDisabled = 'disabled'
            target = ''
            mapHint = ''

            # Worker feedback case.
            # This has a workerId.
            try:
                workerId = req.params['workerId']
                refJson, workJson = mapc.getFeedbackJson(kmlName, workerId)
                instructions = 'Please select one of the overlays and click on a mapped field to see its category and comment labels.'
                commentsVisible = 'none'

            # Standalone case.
            # This has no workerId.
            except:
                workerId = ''
                commentsVisible = 'block'

        # If mapping or training case,
        if len(assignmentId) > 0:
            # Mapping case.
            if len(hitId) > 0:
                k.write("getkml: Mapping request fetched %s kml = %s\n" % (kmlTypeDescr, kmlName))
                k.write("getkml: Mapping request hitId = %s\n" % hitId)
                k.write("getkml: Mapping request assignmentId = %s\n" % assignmentId)
            # Else, training case.
            else:
                k.write("getkml: Training request fetched %s kml = %s\n" % (kmlTypeDescr, kmlName))
                k.write("getkml: Training request assignmentId = %s\n" % assignmentId)
        # Else, worker feedback or standalone cases.
        else:
            if len(workerId) > 0:
                k.write("getkml: Worker feedback request fetched %s kml = %s\n" % (kmlTypeDescr, kmlName))
            else:
                k.write("getkml: Standalone request fetched %s kml = %s\n" % (kmlTypeDescr, kmlName))

        mainText = '''
            <!DOCTYPE html>

                    <title>Labelling Region</title>
                    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
                    <link rel="stylesheet" href="https://fonts.googleapis.com/css?family=Lato:300,400">
                    <link rel="stylesheet" href="/OL/ol.css" type="text/css">
                    <link rel="stylesheet" href="/OL/ol-ext.min.css" type="text/css">
                    <link rel="stylesheet" href="/OL/ol3-layerswitcher.css" type="text/css">
                    <link rel="stylesheet" href="/OL/fontello-88e97ad9/css/fontello.css" type="text/css" />
                    <link rel="stylesheet" href="/OL/controlbar.css" type="text/css">
                    <link rel="stylesheet" href="/OL/showkml.css" type="text/css">

                    <form style='width:100%%;' name='mappingform' action='%(submitTo)s' method='POST' target='%(target)s'>

                            %(instructions)s

                            For comments, problems, or questions:

                            <input type='text'  class='comments' name='comment' size=80 maxlength=2048 %(commentsDisabled)s></input>

                            Hover over the icons in the toolbars below for usage instructions.

                        %(mapHint)s
                        %(csrfToken)s
                        <input type='hidden' name='kmlName' value='%(kmlName)s' />
                        <input type='hidden' name='hitId' value='%(hitId)s' />
                        <input type='hidden' name='assignmentId' value='%(assignmentId)s' />
                        <input type='hidden' name='tryNum' value='%(tryNum)s' />
                        <input type='hidden' name='savedMaps' />
                        <input type='hidden' name='kmlData' />
                    </form>

%(select)s

<textarea id="commentLabel" placeholder="Optional comment" title="Use this box to enter an optional comment"></textarea>

<button id="labelDone" title="Click 'Save Label' to record your field category and optional comment">Save Label</button>

        ''' % {
            'gridJson': gridJson,
            'gridJson2': gridJson2,
            'kmlName': kmlName,
            'hitId': hitId,
            'assignmentId': assignmentId,
            'tryNum': tryNum,
            'resultsAccepted': resultsAccepted,
            'submitTo': submitTo,
            'target': target,
            'instructions': instructions,
            'commentsDisabled': commentsDisabled,
            'commentsVisible': commentsVisible,
            'mapHint': mapHint,
            'kmlMapHeight': kmlMapHeight,
            'refJson':refJson,
            'workJson': workJson,
            'select': select,
            'csrfToken': csrfToken,
            'imageAttributes': image_attributes,
            'instanceid': instanceid,
            'snapTolerance': snapTolerance
        }
        res.text = mainText
    # No KML specified.
    else:
        mainText = '''
            <!DOCTYPE html>

                    <title>Labelling region</title>
                    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">

                    No KML specified in URL.

        '''
        res.body = mainText
        k.write("getkml: No KML specified in URL.\n")
    del mapc
    k.close()
    return res(environ, start_response)


from model_mommy import mommy
from django.urls import reverse
from mydentist.models import Doctor

def test_doctor_only_from_facility(doctor_client):
    mommy.make(Doctor, _quantity=2)
    url = reverse("doctors-list")
    response = doctor_client.get(url)
    data = response.data
    assert data
    assert Doctor.objects.count() == 3
    assert len(data) == 1


import logging

from tg import tmpl_context as c

from ming.odm import ThreadLocalODMSession

from forgewiki import model as WM

log = logging.getLogger(__name__)

def main():
    c.project = None
    pages = WM.Page.query.find({'title': {'$regex': r'\/'}}).all()
    print('Found %s wiki titles containing "/"...' % len(pages))
    for page in pages:
        page.title = page.title.replace('/', '-')
        print('Updated: %s' % page.title)
    ThreadLocalODMSession.flush_all()

if __name__ == '__main__':
    main()


import json
import numpy as np
import os
import reconstruction as rct
from file_reader import roiread

'''
    A colections of points Point() are correlated in a manner that creates
    a perimeter Perimeter(). Every perimeter should be nice, i.e.:
        1) Not self-intersecting;
        2) No overlaping points;
        3) Have a prefered orientation.
    If we can garantee this properties, than we proceed to stitch a colection
    of perimeters in a surface Surface().
'''
with open('main.json', 'r') as settings:
    data = settings.read()
Data = json.loads(data)
FileDir = Data["FileDir"]
OutputDir = Data["OutputDir"]
DisplayConsoleStats = Data["DisplayConsoleStats"]
MeshObjOutput = Data["MeshObjOutput"]

print("Loading files\n\n")
print(FileDir)

def island_init(file_dir,f,subdivision=3):
    arq = roiread(file_dir+"/"+f)
    I = rct.Perimeter(arq)
    I.remove_overlap()
    I.remove_overlap()
    I.fix_intersection()
    I.fix_distance(subdivision=3)
    I.c_clockwise()
    return I

for block in Data["Stitches3D"]:
    for section in block:
        S = rct.Surface()
        for file in block[section]:
            try:
                if isinstance(file,list):
                    I = 0
                    for f in file:
                        I_s = island_init(FileDir,f,3)
                        if I == 0:
                            I = I_s
                        else:
                            I.islands_ensemble(I_s)
                    I.fix_intersection()
                else:
                    I = island_init(FileDir,file,3)
                S.add_island(I)
            except Exception:
                print("Failed to load "+file)

        print("\nBuilding surface: ",section)
        S.build_surface()

        with open("Kamilla_interna_"+section+".obj", "w") as out_file:
            out_file.write(S.surfaceV)
            out_file.write(S.surfaceE)


import urllib.request
from bs4 import BeautifulSoup
import os
import json
import requests
from .utils import get_params

class SniperSpider:
    """Used to get a single inofrmation on a single web page.
    """
    def __init__(self, params={}):
        self.url = None
        self.request = None

        self.http_string = None
        self.expected_value = None

        self.set_params_from_dict(params=params)

    def set_params_from_dict(self, params):
        self.url = params.get("url", "")
        self.request = params.get("request", [])
        self.expected_value = params.get("expected_value", "")

    # ==================================================================

    def shoot(self):
        """Load the web page and get the desired information from it
        Returns:
            str: info from web page
        """
        target = ""
        # Get the web page
        self._get_http_string()
        # Try to access the data in the web page
        try:
            target = self._aim()
        except:
            target = "None"
            print(f"The spider of {self.url} couldn't aim at the target")
        # determining whether the page was parsed and the spider succeeded
        success, parsed = self._pull_triger(target)
        return target, success, parsed

    def _get_http_string(self):
        """call the web page, format it to be readable by beatifulsoup
         and store the result in http_string
        """
        try:
            rep = requests.get(self.url, timeout=0.5)
            self.http_string = rep.content.decode("utf8")
            self.http_string = BeautifulSoup(self.http_string, 'html.parser')
        except:
            self.http_string = "None"
            print(f"Couldn't get {self.url} http string")

    def _aim(self):
        """Read the web page and get the information

        Args:
            request (dictionary): defines the cell information where the data is.

        Returns:
            str: info from web page
        """
        if self.http_string != "None":
            http_sel = self.http_string
            for step in self.request:
                # assign parameters of the parsing
                params = {"name": step["cell_type"]}
                if step.get("class", ""):
                    params["class_"]= step["class"]
                http_sel = http_sel.find(**params)
                # result is the class attribute of the cell if it has been found
                if step.get("test_presence", False) and http_sel:
                    result = step.get("class", "None")
                # result is the content of the cell
                elif step.get("string", False):
                    result = http_sel.string
                else:
                    result = "None"
        else:
            result = "None"
        return result

    def _pull_triger(self, parsing_result):
        """indicate wheter the page was parsed and the spider found the value expected
        Args: parsing_result (string): the value returned by the spider
        Returns: success, parsed (bool, bool)
        """
        if parsing_result == self.expected_value:
            success = True
            parsed = True
        elif parsing_result == "None":
            success = False
            parsed = False
        else:
            success = False
            parsed = True

        return success, parsed

if __name__ == "__main__":
    params = {
        "url": "",
        "request": {
            "cell_type" : "div",
            "class" : "product-page-description col-flex-lg-5 col-flex-sm-12",
            "string" : False
        }
    }
    params = get_params("spider_data")["processor_topachat"]
    #print(params)
    sniper = SniperSpider(params)

    result, success, parsed = sniper.shoot()

    print(result, success, parsed)


import csv
from datetime import datetime
from convert import convert
from matplotlib import pyplot as plt

def main(fl):
    filenames=fl  #16.1.1
    with open(filenames) as f:
        reader = csv.reader(f)
        header_row=next(reader)
        print(header_row)

        datas,highs,lows=[],[],[]
        for row in reader:
            try:
                current_date=datetime.strptime(row[0],"%Y-%m-%d")
                high=convert(float(row[1]))
                low= convert(float(row[3]))

            except ValueError:
                print(current_date,'missing data')

            else:
                datas.append(current_date)
                highs.append(high)
                lows.append(low)

    title=[]
    title.append(datas)
    title.append(highs)
    title.append(lows)
    return title

def high_lows(weather,weather_1=''):
    datas,highs,lows=weather
    if weather_1 != 0:
        datas_1,highs_1,lows_1=weather_1

    fig=plt.figure(dpi=128,figsize=(10,6))
    plt.plot(datas,highs,c='red',alpha=0.5)
    plt.plot(datas,lows,c='blue',alpha=0.5)
    plt.fill_between(datas,highs,lows,facecolor='blue',alpha=0.1)

    if weather_1 != 0:
        plt.plot(datas_1,highs_1,c='#C71585',alpha=0.5)
        plt.plot(datas_1,lows_1,c='#8A2BE2',alpha=0.5)
        plt.fill_between(datas_1,highs_1,lows_1,facecolor='#FFA500',alpha=0.1)

    plt.title("Daily tempretures, 2014",fontsize=24)
    plt.xlabel('',fontsize=16)
    fig.autofmt_xdate()
    plt.ylabel("Temperature(C)",fontsize=16)
    plt.tick_params(axis='both',which='major',labelsize=16)

    plt.show()

weather=main('sitka_weather_07-2014.csv')
weather_1=main('death_valley_2014.csv')
high_lows(weather,weather_1)


import pandas as pd
import os
import matplotlib.pyplot as plt
import numpy as np

def plot_df(df_results, out_path=None):
    f1s = df_results['F1-score'].values
    plt.figure(figsize=(10, 10))
    plt.hist(f1s, bins='auto')
    plt.title("F1-distribution")
    plt.ylabel('Absolute frequency')
    plt.xlabel('F1-score')
    if out_path:
        plt.savefig(out_path)
        plt.close()
    else:
        plt.show()
        plt.close()

def make_results_table2(df_results, pipeline_name, round_dec=1):
    df_results = df_results[["Precision", "Recall", "F1-score"]]
    percentiles = np.transpose(np.quantile(df_results, q=[0.025, 0.5, 0.975], axis=0))
    percentiles = np.round(percentiles, 4) * 100
    precision = "{0} ({1},{2})".format(str(np.round(percentiles[0][1], round_dec)), str(np.round(percentiles[0][0], round_dec)),
                                       str(np.round(percentiles[0][2], round_dec)))
    recall = "{0} ({1},{2})".format(str(np.round(percentiles[1][1], round_dec)), str(np.round(percentiles[1][0], round_dec)),
                                    str(np.round(percentiles[1][2], round_dec)))
    f1 = "{0} ({1},{2})".format(str(np.round(percentiles[2][1], round_dec)), str(np.round(percentiles[2][0], round_dec)),
                                str(np.round(percentiles[2][2], round_dec)))

    iqv = np.round(percentiles[2][2], round_dec) - np.round(percentiles[2][0], round_dec)
    data = [[pipeline_name, precision, recall, f1, iqv]]
    out_dataframe = pd.DataFrame(data, columns=['Pipeline', 'Precision', 'Recall', 'F1-score', 'IQV'])
    return out_dataframe

def extract_f1(inp_text):
    return float(inp_text.split("(")[0])

def rename(inp_word):
    mapper = dict(res_title='Title', res_authors='Authors', res_journal='Journal', res_keywords='Keywords',
                  res_abstract='Abstract', res_chemical='Chemicals', res_mesh='MeSH', res_pub_type='Pub. Type',
                  res_affiliations='Affiliations', res_all='All fields', res_optimal='Opt. Fields',
                  res_optimal_idf='Opt. Fields IDF', res_unigrams='Unigrams', res_bigrams='Bigrams',
                  res_trigrams='Trigrams', res_specter='SPECTER', res_biobert_avg='BioBERT\nmean pooling',
                  res_biobert_all='BioBERT mean\n+ min&max pooling', res_biobert_all_bow='Unigrams\n+ BioBERT mean\n+ '
                                                                                'min&max pooling',
                  res_biobert_avg_bow='Unigrams\n+ BioBERT\nmean pooling')
    if inp_word in mapper.keys():
        return mapper[inp_word]
    else:
        return inp_word

def get_all_results(inp_result_files, input_dir, output_dir, convert_latex):
    all_results = []
    all_for_boxplot = []
    all_for_boxplot_names = []

    if 'res_unigrams.csv' in inp_result_files:
        inp_result_files = inp_result_files[-1:] + inp_result_files[:-1]

    print(inp_result_files)

    for inp_result in inp_result_files:
        instance_df = pd.read_csv(os.path.join(input_dir, inp_result))
        all_for_boxplot.append(instance_df["F1-score"].values)
        instance_name = inp_result.split(".")[0]
        all_for_boxplot_names.append(instance_name)
        results = make_results_table2(instance_df, instance_name)
        all_results.append(results)

    all_results_ready = pd.concat(all_results)
    all_results_ready['F1'] = all_results_ready['F1-score'].apply(extract_f1)
    all_results_ready = all_results_ready.sort_values(by="F1-score", ascending=False)
    print(all_results_ready)
    if convert_latex:
        print(all_results_ready.to_latex(index=False))

    all_results_ready.to_csv(os.path.join(output_dir, "all_results.csv"))

    idx_medians_sorted = np.asarray([np.median(x) for x in all_for_boxplot]).argsort()
    # idx_medians_sorted = range(0,len(idx_medians_sorted))
    all_for_boxplot = [all_for_boxplot[i] * 100 for i in idx_medians_sorted]
    all_for_boxplot_names = [rename(all_for_boxplot_names[i]) for i in idx_medians_sorted]

    fig7, ax7 = plt.subplots()
    fig7.set_figheight(40)
    fig7.set_figwidth(10)

    plt.ylim(55, 90)
    ax7.boxplot(all_for_boxplot, labels=all_for_boxplot_names, whis=[2.5, 97.5])
    plt.ylabel('F1-score (%)', fontsize=16)
    plt.xticks(rotation=65, fontsize=12)
    plt.yticks(fontsize=12)
    plt.gcf().subplots_adjust(bottom=0.2)
    plt.show()
    plt.close()


import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision.transforms import transforms

# 设置随机种子，以便结果可重复
torch.manual_seed(42)

# 检查是否有可用的GPU设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小为224x224像素
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化图像
])

# 加载训练集和验证集
train_data = torchvision.datasets.ImageFolder('path_to_train_data', transform=transform)
val_data = torchvision.datasets.ImageFolder('path_to_validation_data', transform=transform)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=32, shuffle=False)

# 定义CNN模型
class FoodClassifier(nn.Module):
    def __init__(self):
        super(FoodClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(256 * 28 * 28, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 10)  # 假设有10个食物类别
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

# 创建模型实例
model = FoodClassifier().to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    train_correct = 0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        _, predicted = torch.max(outputs.data, 1)
        train_correct += (predicted == labels).sum().item()
        train_loss += loss.item() * images.size(0)

    train_accuracy = train_correct / len(train_data)
    train_loss = train_loss / len(train_data)

    # 在验证集上评估模型
    model.eval()
    val_loss = 0.0
    val_correct = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            _, predicted = torch.max(outputs.data, 1)
            val_correct += (predicted == labels).sum().item()
            val_loss += loss.item() * images.size(0)

    val_accuracy = val_correct / len(val_data)
    val_loss = val_loss / len(val_data)

    print(f"Epoch [{epoch+1}/{num_epochs}], "
          f"Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}, "
          f"Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}")

# 保存模型
torch.save(model.state_dict(), 'food_classifier.pth')


# -*- coding: utf-8 -*-
from zerotk.zops import call_main, Console

import click

click.disable_unicode_literals_warning = True

@click.group('code-standards')
def main():
    pass

@main.command()
@click.argument('paths', nargs=-1)
@click.pass_context
def apply(ctx, paths):
    Console.title('zops code-standards apply')
    for i_path in paths:
        ctx.invoke(isort, path=i_path)
        ctx.invoke(autoflake, path=i_path)
        ctx.invoke(autopep8, path=i_path)

@main.command()
@click.argument('path')
def isort(path):
    import os
    from isort.main import main as isort_main

    Console.title('zops code-standards isort')
    isort_args = [
        #'--verbose',
        '-m3',
        '-l120',
        #'--order-by-type',
        #'--combine-as',
        #'--force-alphabetical-sort-within-sections',
        #'--force-alphabetical-sort',
        #'--force-sort-within-sections',
        #'--lines-between-types=1',
        #'--atomic',
        #'-y',
        path,
    ]
    if os.path.isdir(path):
        isort_args.insert(0, '--recursive')
    Console.execution('isort', *isort_args)
    call_main(isort_main, *isort_args)

@main.command()
@click.argument('path')
def autoflake(path):
    import os
    from autoflake import main as autoflake_main

    Console.title('zops code-standards autoflake')
    autoflake_args = [
        '-i',
        '--remove-all-unused-imports',
        path
    ]
    if os.path.isdir(path):
        autoflake_args.insert(0, '--recursive')
    Console.execution('autoflake', *autoflake_args)
    call_main(autoflake_main, *autoflake_args)

@main.command()
@click.argument('path')
def autopep8(path):
    import os
    from autopep8 import main as autopep8_main

    Console.title('zops code-standards autopep8')
    autopep8_args = [
        '-a',
        '-i',
        '--max-line-length=120',
        path
    ]
    if os.path.isdir(path):
        autopep8_args.insert(0, '--recursive')
    Console.execution('autopep8', *autopep8_args)
    call_main(autopep8_main, *autopep8_args)


#!/usr/bin/python
# -*- coding: utf-8 -*-

# '''
# This values moved here from *.cfg for ability of changing from elsewhere
# '''

from multiprocessing import Lock

flags = {"lock" : Lock(), # для атомарного доступа к флагам
         "atm_mode" : False, # Режим настройки дрессированной обезьяной
         "debug" : False, # включает отладочный режим
         # "debug" : True, # включает отладочный режим
         "mouse_capture" : False, # включает захват координат курсора в качестве входных данных
         # "mouse_capture" : True, # включает захват координат курсора в качестве входных данных
         "help_flag" : False,# если взведен -- показывать экран помощи
         "logreplay_flag" : False, # если взведен -- вывести на экран значок play
         "connection_interrupted" : False, # связь с мпиовд прервалась
         }


import pdb
import random
import torch
from torch.optim import Adam
from torch.nn import functional as F
from tqdm import tqdm
from model import Model
from tokenizing import create_vocab

def main():

  print("building vocab")

  with open("data/all_notes.txt", "r") as f:
    full_corpus = f.read()

  token_map, char_map = create_vocab(full_corpus)

  print("loading splits")

  with open("data/train_split.txt", "r") as f:
    train_corpus = f.read()

  with open("data/test_split.txt", "r") as f:
    test_corpus = f.read()

  device = "cpu"  #"cuda" # cpu

  context_size = 8 #16 #64
  hidden_dim = 16 #128
  vocab_size = len(token_map.keys())

  print("initializing model")

  model = Model(vocab_size, context_size=context_size, hidden_dim=hidden_dim, layers=2).to(device)

  iterations = 100
  steps_per_iter = 20000 #10000
  test_steps = 100
  batch_size = 64

  opt = Adam(model.parameters())

  with torch.no_grad():
    init_inp, init_out = get_batch(test_corpus, token_map, batch_size, context_size, device)
    print(f"initial sample: \n{generate_sample(model, init_inp[0:1], char_map)}")

  print("begin training")

  for i in range(iterations):
    mean_train_loss = 0
    for e in range(steps_per_iter):
      input_t, output_t = get_batch(train_corpus, token_map, batch_size, context_size, device)
      opt.zero_grad()
      logits = model(input_t)
      # probs = F.softmax(logits, dim=1)
      #if i == 5:
      #  pdb.set_trace()
      loss = F.cross_entropy(logits, output_t)
      mean_train_loss += loss.cpu().item()
      loss.backward()
      opt.step()

    mean_train_loss /= steps_per_iter
    print(f"train loss: {mean_train_loss}")
    eval_in, eval_out = get_batch(test_corpus, token_map, batch_size, context_size, device)
    with torch.no_grad():
      mean_test_loss = 0
      for tb in range(test_steps):
        input_t, output_t = get_batch(test_corpus, token_map, batch_size, context_size, device)
        logits = model(input_t)
        loss = F.cross_entropy(logits, output_t)
        mean_test_loss += loss.cpu().item()
      mean_test_loss /= test_steps
      print(f"test loss: {mean_test_loss}")
      sampled_text = generate_sample(model, eval_in[0:1], char_map)
      print(f"{sampled_text}")

def generate_sample(model, curr_context, char_map):
  raw_chars = [char_map[t.item()] for t in curr_context[0][:16]]
  new_text = ""
  for c in raw_chars:
    new_text += c
  new_text += "|generation-start|"
  for i in range(100):
    logits = model(curr_context)
    #probs = F.softmax(logits, dim=1)[0]
    new_token = torch.distributions.Categorical(logits=logits).sample()
    new_text += char_map[new_token.item()]
    curr_context[0] = curr_context[0].roll(-1)
    curr_context[0][-1] = new_token
    #pdb.set_trace()

  return new_text + "|generation-end"

def get_batch(data, token_map, batch_size, context_size, device):
  max_idx = len(data) - context_size - 1
  offsets = [random.randint(0, max_idx) for _ in range(batch_size)]
  input_tokens = torch.tensor([[token_map[c] for c in data[off:off+context_size]] for off in offsets], device=device)
  #print(f"input chars : {input_tokens}")
  output_tokens = torch.tensor([token_map[data[off+context_size]] for off in offsets], device=device)
  #print(f"next chars : {output_tokens}")
  return input_tokens, output_tokens

if __name__ == "__main__":
  main()


import urllib.request
import csv
import sys
import io
import pandas as pd
from functools import reduce

url = "https://data.cityofnewyork.us/Public-Safety/FDNY-Line-Of-Duty-Deaths/32y8-s55cs/data.csv" \
    "?accessType=DOWNLOAD"

try:
    fileFromUrl = urllib.request.urlopen(url)
except urllib.error.URLError as error:
    print("urllib.error.URLError", error)
    sys.exit(1)

sequenceOfBytes = fileFromUrl.read() #Read whole file into one big sequenceOfBytes.
fileFromUrl.close()

try:
    s = sequenceOfBytes.decode("utf-8")    #s is a string
except UnicodeError as unicodeError:
    print(unicodeError)
    sys.exit(1)

fileFromString = io.StringIO(s)
df = pd.read_csv(fileFromString, dtype = {'str'})  #reads in fileFromString as DataFrame
fileFromString.close()

deaths=[
    "Rank",
    "Name",
    "Unit",
    "Date"
    ]

deaths.drop(columns=deaths, inplace=True)

print("FDNY NY DEATHS")
print()
print("Rank")
print()
print("Name")
print()
print("Unit")
print()
print("Date")

sys.exit(0)