#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Command line interface for ms3. This module includes the argument parsers for the various commands and has one function
<command>_cmd() per <command> which takes the parsed arguments namespace and dispatches them to the appropriate function
in the ms3.operations module, using a couple of helper functions to that aim.
"""

import argparse
import os
import re
from collections import defaultdict
from typing import List, Optional

from ms3 import Parse, compute_path_from_file, make_coloring_reports_and_warnings
from ms3._version import __version__
from ms3.logger import get_logger, inspect_loggers
from ms3.operations import (
    check,
    compare,
    extract,
    insert_labels_into_score,
    store_scores,
    transform_to_package,
    transform_to_resources,
    update,
)
from ms3.utils import (
    capture_parse_logs,
    convert_folder,
    resolve_dir,
    write_warnings_to_file,
)

__author__ = "johentsch"
__copyright__ = "École Polytechnique Fédérale de Lausanne"
__license__ = "gpl3"

def gather_extract_params(args) -> List[str]:
    params = [
        name
        for name, arg in zip(
            (
                "measures",
                "notes",
                "rests",
                "labels",
                "expanded",
                "form_labels",
                "events",
                "chords",
            ),
            (
                args.measures,
                args.notes,
                args.rests,
                args.labels,
                args.expanded,
                args.form_labels,
                args.events,
                args.chords,
            ),
        )
        if arg
    ]
    if args.metadata is not None and not (args.metadata == False):  # noqa
        params.append("metadata")
    return params

def add_cmd(args, parse_obj: Optional[Parse] = None) -> Parse:
    logger = get_logger("ms3.add", level=args.level)
    if parse_obj is None:
        p = make_parse_obj(args)
    else:
        p = parse_obj

    insert_labels_into_score(
        p,
        facet=args.use,
        ask_for_input=args.ask,
        replace=args.replace,
    )
    corpus2paths = store_scores(
        p, root_dir=args.out, folder=".", simulate=args.test, suffix=args.suffix
    )
    changed = sum(map(len, corpus2paths.values()))
    logger.info(
        f"Operation resulted in {changed} annotated score{'s' if changed != 1 else ''}."
    )

def check_cmd(args, parse_obj: Optional[Parse] = None) -> Parse:
    if parse_obj is None:
        p = make_parse_obj(args)
    else:
        p = parse_obj
    _ = check(
        p,
        ignore_labels=args.ignore_labels,
        ignore_scores=args.ignore_scores,
        assertion=args.fail,
        ignore_metronome=args.ignore_metronome,
    )
    return p

def compare_cmd(args, parse_obj=None, output: bool = True, logger=None) -> None:
    if logger is None:
        logger = get_logger("ms3.compare", level=args.level)
    if parse_obj is None:
        p = make_parse_obj(args)
    else:
        p = parse_obj
    revision = None if args.compare == "" else args.compare
    try:
        n_changed, n_unchanged = compare(
            p,
            facet=args.use,
            ask=args.ask,
            revision_specifier=revision,
            flip=args.flip,
            logger=p.logger,
        )
    except TypeError:
        logger.error(
            "No files to compare. Try adding the flag -a to parse files even if they are not included in a "
            "metdata.tsv file."
        )
        return
    logger.debug(
        f"{n_changed} files changed labels during comparison, {n_unchanged} didn't."
    )
    if output:
        corpus2paths = store_scores(
            p, root_dir=args.out, simulate=args.test, suffix=args.suffix
        )
        changed = sum(map(len, corpus2paths.values()))
        logger.info(
            f"Operation resulted in {changed} comparison file{'s' if changed != 1 else ''}."
        )
    else:
        logger.info(
            f"Operation resulted in colored labels within {n_changed} score{'s' if n_changed != 1 else ''}."
        )

def convert_cmd(args):
    # assert target[:len(
    #    dir)] != dir, "TARGET_DIR cannot be identical with nor a subfolder of DIR.\nDIR:        " + dir +
    #    '\nTARGET_DIR: ' + target
    update_logger = get_logger("ms3.convert", level=args.level)
    for argument_name, argument, default in (
        ("-a/--all", args.all, False),
        ("-e/--exclude", args.exclude, None),
        ("-f/--folders", args.folders, None),
        ("--reviewed", args.reviewed, False),
        ("-t/--test", args.test, False),
        ("-v/--verbose", args.verbose, False),
    ):
        if argument != default:
            update_logger.info(
                f"Argument '{argument_name}' is currently being ignored."
            )
    out_dir = os.getcwd() if args.out is None else resolve_dir(args.out)
    ms = "auto" if args.musescore is None else args.musescore
    convert_folder(
        directory=args.dir,
        file_paths=args.files,
        target_dir=out_dir,
        extensions=args.extensions,
        target_extension=args.format,
        regex=args.include,
        suffix=args.suffix,
        recursive=not args.nonrecursive,
        ms=ms,
        overwrite=args.safe,
        parallel=not args.iterative,
        logger=update_logger,
    )

def empty(args, parse_obj: Optional[Parse] = None):
    if parse_obj is None:
        p = make_parse_obj(args, parse_scores=True)
    else:
        p = parse_obj
    p.detach_labels()
    p.store_parsed_scores(
        only_changed=True,
        root_dir=args.out,
        suffix=args.suffix,
        overwrite=True,
    )

def extract_cmd(args, parse_obj: Optional[Parse] = None):
    if parse_obj is None:
        p = make_parse_obj(args)
    else:
        p = parse_obj
    params = gather_extract_params(args)
    if len(params) == 0:
        print(
            "In order to extract DataFrames, pass at least one of the following arguments: -M (measures), -N (notes), "
            "-R (rests), -L (labels), -X (expanded), -F (form_labels), -E (events), -C (chords), -D (metadata)"
        )
        return
    extract(
        p,
        root_dir=args.out,
        notes_folder=args.notes,
        labels_folder=args.labels,
        measures_folder=args.measures,
        rests_folder=args.rests,
        events_folder=args.events,
        chords_folder=args.chords,
        expanded_folder=args.expanded,
        form_labels_folder=args.form_labels,
        metadata_suffix=args.metadata,
        simulate=args.test,
        parallel=not args.iterative,
        unfold=args.unfold,
        interval_index=args.interval_index,
        corpuswise=args.corpuswise,
    )

def metadata(args, parse_obj: Optional[Parse] = None):
    """Update MSCX files with changes made in metadata.tsv (created via ms3 extract -D). In particular,
    add the values from (new?) columns to the corresponding fields in the MuseScore files' "Score info".
    """
    if parse_obj is None:
        p = make_parse_obj(args)
    else:
        p = parse_obj
    print("PARSING SCORES ONE AFTER THE OTHER...")
    p.parse_scores(parallel=False)
    modified_score_files = p.update_score_metadata_from_tsv(
        write_empty_values=args.empty,
        remove_unused_fields=args.remove,
        write_text_fields=args.prelims,
        update_instrumentation=args.instrumentation,
    )
    if len(modified_score_files) == 0:
        print("Nothing to update.")
        return
    corpus2paths = store_scores(
        p,
        only_changed=True,
        root_dir=args.out,
        folder=".",
        simulate=args.test,
        suffix=args.suffix,
    )
    changed = sum(map(len, corpus2paths.values()))
    print(
        f"Operation resulted in {changed} updated score{'s' if changed != 1 else ''}."
    )

# def repair(args):
#     print("Sorry, the command has not been implemented yet.")
#     print(args.dir)

def transform_cmd(args):
    params = gather_extract_params(args)
    if len(params) == 0:
        print(
            "Pass at least one of the following arguments: -M (measures), -N (notes), -R (rests), -L (labels), "
            "-X (expanded), -F (form_labels), -E (events), -C (chords), -D (metadata)"
        )
        return
    parse_obj = make_parse_obj(args, parse_tsv=True, facets=params)
    filename = os.path.basename(args.dir)
    func = transform_to_resources if args.resources else transform_to_package
    # noinspection PyTypeChecker
    func(
        ms3_object=parse_obj,
        facets=params,
        filename=filename,
        output_folder=args.out,
        choose="auto",
        interval_index=args.interval_index,
        unfold=args.unfold,
        test=args.test,
        zipped=not args.uncompressed,
        overwrite=args.safe,
        log_level=args.level,
    )

def update_cmd(args, parse_obj: Optional[Parse] = None):
    if parse_obj is None:
        p = make_parse_obj(args, parse_scores=True)
    else:
        p = parse_obj
    changed_paths = update(
        p,
        root_dir=args.out,
        suffix="" if args.suffix is None else args.suffix,
        overwrite=True,
        staff=int(args.staff),
        harmony_layer=args.type,
        above=args.above,
        safe=args.safe,
    )
    if len(changed_paths) > 0:
        print(f"Paths of scores with updated labels:\n{changed_paths}")
    else:
        print("Nothing to do.")

def check_and_create(d):
    """Turn input into an existing, absolute directory path."""
    if not os.path.isdir(d):
        d = resolve_dir(os.path.join(os.getcwd(), d))
        if not os.path.isdir(d):
            if input(d + " does not exist. Create? (y|n)") == "y":
                os.mkdir(d)
            else:
                raise argparse.ArgumentTypeError(
                    d + " needs to be an existing directory"
                )
    return resolve_dir(d)

def check_dir(d):
    if not os.path.isdir(d):
        d = resolve_dir(os.path.join(os.getcwd(), d))
        if not os.path.isdir(d):
            raise argparse.ArgumentTypeError(d + " needs to be an existing directory")
    return resolve_dir(d)

def review_cmd(
    args,
    parse_obj: Optional[Parse] = None,
) -> None:
    """This command combines the functionalities check-extract-compare and additionally colors non-chord tones."""
    logger = get_logger("ms3.review", level=args.level)
    if parse_obj is None:
        p = make_parse_obj(args)
    else:
        p = parse_obj
    accumulated_warnings = []

    # call ms3 check, collecting all warnings issued while parsing
    if args.ignore_scores + args.ignore_labels < 2:
        what = "" if args.ignore_scores else "SCORES"
        if args.ignore_labels:
            what += "LABELS" if what == "" else "AND LABELS"
        print(f"CHECKING {what}...")
        warning_strings = check(
            p,
            ignore_labels=args.ignore_labels,
            ignore_scores=args.ignore_scores,
            assertion=False,
            parallel=False,
            ignore_metronome=args.ignore_metronome,
        )
        accumulated_warnings.extend(warning_strings)
        test_passes = len(warning_strings) == 0
        p.parse_tsv()
    else:
        test_passes = True
        p.parse(parallel=False)
    if p.n_parsed_scores == 0:
        msg = "NO SCORES PARSED, NOTHING TO DO."
        if not args.all:
            msg += (
                "\nI was disregarding all files whose file names are not listed in the column 'piece' of a "
                "'metadata.tsv' file. "
                "Add -a if you want me to include all scores."
            )
        print(msg)
        return

    # call ms3 extract
    params = gather_extract_params(args)
    if len(params) > 0:
        print("EXTRACTING FACETS...")
        extract_cmd(args, p)

    # color out-of-label tones
    print("COLORING OUT-OF-LABEL TONES...")
    with capture_parse_logs(p.logger) as captured_warnings:
        test_passes = (
            make_coloring_reports_and_warnings(
                p, out_dir=args.out, threshold=args.threshold
            )
            and test_passes
        )
        accumulated_warnings.extend(captured_warnings.content_list)

    # attribute warnings to pieces based on logger names
    logger2pieceID = {
        corpus.logger_names[piece]: (c, piece)
        for c, corpus in p.corpus_objects.items()
        for piece in corpus.keys()
    }
    piece2warnings = defaultdict(list)
    for warning_strings in accumulated_warnings:
        warning_lines = warning_strings.splitlines()
        first_line = warning_lines[0]
        match = re.search(r"(ms3\.Parse\..+) --", first_line)
        if match is None:
            logger.warning(
                f"This warning contains no ms3 logger name, skipping: {warning_strings}"
            )
            continue
        warning_lines[0] = first_line[
            : match.end(1)
        ]  # cut off the warning's header everything following the logger
        # name because paths to source code are system-dependent
        pieceID = logger2pieceID[match.group(1)]
        piece2warnings[pieceID].append("\n".join(warning_lines))

    # write all warnings to piece-specific warnings files and remove existing files where no warnings were captured
    for pieceID, score_files in p.get_files(
        "scores", unparsed=False, flat=True, include_empty=False
    ).items():
        file = score_files[0]
        warnings_path = compute_path_from_file(
            file, root_dir=args.out, folder="reviewed"
        )
        warnings_file = os.path.join(
            warnings_path, file.piece + file.suffix + ".warnings"
        )
        write_warnings_to_file(
            warnings_file=warnings_file, warnings=piece2warnings[pieceID], logger=logger
        )

    # call ms3 compare
    if args.compare is not None:
        revision = None if args.compare == "" else args.compare
        commit_str = "" if revision is None else f" @ {revision}"
        expanded = "EXPANDED " if args.use == "expanded" else ""
        print(
            f"COMPARING CURRENT {expanded}LABELS WITH PREVIOUS ONES FROM THE TSV FILES{commit_str}..."
        )
        n_changed, n_unchanged = compare(
            p,
            facet=args.use,
            ask=args.ask,
            revision_specifier=revision,
            flip=args.flip,
            logger=logger,
        )
        logger.debug(
            f"{n_changed} files changed labels during comparison, {n_unchanged} didn't."
        )
    corpus2paths = store_scores(p, root_dir=args.out, simulate=args.test)
    changed = sum(map(len, corpus2paths.values()))
    logger.info(
        f"Operation resulted in {changed} review file{'s' if changed != 1 else ''}."
    )
    if test_passes:
        logger.info("Parsed scores passed all tests.")
    else:
        msg = "Not all tests have passed. Please check the relevant warnings in the generated .warnings files."
        if args.fail:
            assert test_passes, msg
        else:
            logger.info(msg)
    out_dir = "." if args.out is None else args.out
    warnings_file = os.path.join(out_dir, "warnings.log")
    if os.path.isfile(warnings_file):
        logger.info(f"Removing deprecated warnings file: {warnings_file}")
        os.remove(warnings_file)

def make_parse_obj(args, parse_scores=False, parse_tsv=False, facets=None):
    labels_cfg = {}
    if hasattr(args, "positioning"):
        labels_cfg["positioning"] = args.positioning
    if hasattr(args, "raw"):
        labels_cfg["decode"] = args.raw

    logger_cfg = {
        "level": args.level,
        "path": args.log,
    }
    ms = args.musescore if hasattr(args, "musescore") else None
    file_re_str = None if args.include is None else f"'{args.include}'"
    folder_re_str = None if args.folders is None else f"'{args.folders}'"
    exclude_re_str = None if args.exclude is None else f"'{args.exclude}'"
    ms_str = None if ms is None else f"'{ms}'"
    print(
        f"""CREATING PARSE OBJECT WITH THE FOLLOWING PARAMETERS:
Parse('{args.dir}',
     recursive={not args.nonrecursive},
     only_metadata_pieces={not args.all},
     include_convertible={ms is not None},
     exclude_review={not args.reviewed},
     file_re={file_re_str},
     folder_re={folder_re_str},
     exclude_re={exclude_re_str},
     file_paths={args.files},
     labels_cfg={labels_cfg},
     ms={ms_str},
     **{logger_cfg})
"""
    )
    parse_obj = Parse(
        args.dir,
        recursive=not args.nonrecursive,
        only_metadata_pieces=not args.all,
        include_convertible=ms is not None,
        exclude_review=not args.reviewed,
        file_re=args.include,
        folder_re=args.folders,
        exclude_re=args.exclude,
        file_paths=args.files,
        labels_cfg=labels_cfg,
        ms=ms,
        **logger_cfg,
    )
    if facets is not None:
        facets = [f"{f}$" for f in facets if f != "metadata"]
        if len(facets) > 0:
            parse_obj.view.include("facets", *facets)
    if parse_scores:
        mode = "ONE AFTER THE OTHER" if args.iterative else "IN PARALLEL"
        print(f"PARSING SCORES {mode}...")
        parse_obj.parse_scores(parallel=not args.iterative)
    if parse_tsv and facets != []:
        print("PARSING TSV FILES...")
        parse_obj.parse_tsv()
    info_str = parse_obj.info(show_discarded=args.verbose, return_str=True).replace(
        "\n", "\n\t"
    )
    print(f"RESULTING PARSE OBJECT:\n\t{info_str}")
    if args.verbose:
        print(inspect_loggers())
    return parse_obj

def get_arg_parser():
    # reusable argument sets
    parse_args = argparse.ArgumentParser(add_help=False)
    parse_args.add_argument(
        "-d",
        "--dir",
        metavar="DIR",
        default=os.getcwd(),
        type=check_dir,
        help="Folder(s) that will be scanned for input files. Defaults to current working directory if no individual "
        "files are passed via -f.",
    )
    parse_args.add_argument(
        "-o",
        "--out",
        metavar="OUT_DIR",
        type=check_and_create,
        help="Output directory.",
    )
    parse_args.add_argument(
        "-n",
        "--nonrecursive",
        action="store_true",
        help="Treat DIR as single corpus even if it contains corpus directories itself.",
    )
    parse_args.add_argument(
        "-a",
        "--all",
        action="store_true",
        help="By default, only files listed in the 'piece' column of a 'metadata.tsv' file are parsed. With "
        "this option, all files will be parsed.",
    )
    parse_args.add_argument(
        "-i",
        "--include",
        metavar="REGEX",
        help="Select only files whose names include this string or regular expression.",
    )
    parse_args.add_argument(
        "-e",
        "--exclude",
        metavar="REGEX",
        help="Any files or folders (and their subfolders) including this regex will be disregarded."
        "By default, files including '_reviewed' or starting with . or _ or 'concatenated' are excluded.",
    )
    parse_args.add_argument(
        "-f",
        "--folders",
        metavar="REGEX",
        help="Select only folders whose names include this string or regular expression.",
    )
    parse_args.add_argument(
        "-m",
        "--musescore",
        metavar="PATH",
        nargs="?",
        const="auto",
        help="""Command or path of your MuseScore 3 executable. -m by itself will set 'auto' (attempt to use standard
        path for your system).
        Other shortcuts are -m win, -m mac, and -m mscore (for Linux).""",
    )
    parse_args.add_argument(
        "--reviewed",
        action="store_true",
        help="By default, review files and folder are excluded from parsing. With this option, "
        "they will be included, too.",
    )
    parse_args.add_argument(
        "--files",
        metavar="PATHs",
        nargs="+",
        help="(Deprecated) The paths are expected to be within DIR. They will be converted into a view "
        "that includes only the indicated files. This is equivalent to specifying the file names as "
        "a regex via --include (assuming that file names are unique amongst corpora.",
    )
    parse_args.add_argument(
        "--iterative",
        action="store_true",
        help="Do not use all available CPU cores in parallel to speed up batch jobs.",
    )
    parse_args.add_argument(
        "-l",
        "--level",
        metavar="{c, e, w, i, d}",
        default="i",
        help="Choose how many log messages you want to see: c (none), e, w, i, d (maximum)",
    )
    parse_args.add_argument(
        "--log",
        nargs="?",
        const=".",
        help="Can be a file path or directory path. Relative paths are interpreted relative to the current directory.",
    )
    parse_args.add_argument(
        "-t", "--test", action="store_true", help="No data is written to disk."
    )
    parse_args.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Show more output such as files discarded from parsing.",
    )

    check_args = argparse.ArgumentParser(add_help=False)
    check_args.add_argument(
        "--ignore_scores",
        action="store_true",
        help="Don't check scores for encoding errors.",
    )
    check_args.add_argument(
        "--ignore_labels",
        action="store_true",
        help="Don't check DCML labels for syntactic correctness.",
    )
    check_args.add_argument(
        "--fail",
        action="store_true",
        help="If you pass this argument the process will deliberately fail with an AssertionError "
        "when there are any mistakes.",
    )
    check_args.add_argument(
        "--ignore_metronome",
        action="store_true",
        help="Pass this flag if you want the check to pass (not fail) even if there is a warning about a missing "
        "metronome mark in the first bar of the score.",
    )

    compare_args = argparse.ArgumentParser(add_help=False)
    compare_args.add_argument(
        "--flip",
        action="store_true",
        help="Pass this flag to treat the annotation tables as if updating the scores instead of the other way around, "
        "effectively resulting in a swap of the colors in the output files.",
    )
    compare_args.add_argument(
        "--safe", action="store_false", help="Don't overwrite existing files."
    )

    extract_args = argparse.ArgumentParser(add_help=False)
    extract_args.add_argument(
        "-M",
        "--measures",
        metavar="folder",
        nargs="?",
        const="../measures",
        help="Folder where to store TSV files with measure information needed for tasks such as unfolding repetitions.",
    )
    extract_args.add_argument(
        "-N",
        "--notes",
        metavar="folder",
        nargs="?",
        const="../notes",
        help="Folder where to store TSV files with information on all notes.",
    )
    extract_args.add_argument(
        "-R",
        "--rests",
        metavar="folder",
        nargs="?",
        const="../rests",
        help="Folder where to store TSV files with information on all rests.",
    )
    extract_args.add_argument(
        "-L",
        "--labels",
        metavar="folder",
        nargs="?",
        const="../labels",
        help="Folder where to store TSV files with information on all annotation labels.",
    )
    extract_args.add_argument(
        "-X",
        "--expanded",
        metavar="folder",
        nargs="?",
        const="../harmonies",
        help="Folder where to store TSV files with expanded DCML labels.",
    )
    extract_args.add_argument(
        "-F",
        "--form_labels",
        metavar="folder",
        nargs="?",
        const="../form_labels",
        help="Folder where to store TSV files with all form labels.",
    )
    extract_args.add_argument(
        "-E",
        "--events",
        metavar="folder",
        nargs="?",
        const="../events",
        help="Folder where to store TSV files with all events (chords, rests, articulation, etc.) without further "
        "processing.",
    )
    extract_args.add_argument(
        "-C",
        "--chords",
        metavar="folder",
        nargs="?",
        const="../chords",
        help="Folder where to store TSV files with <chord> tags, "
        "i.e. groups of notes in the same voice with identical onset and duration. "
        "The tables include lyrics, dynamics, articulation, staff- and system texts, tempo marking, spanners, "
        "and thoroughbass figures.",
    )
    extract_args.add_argument(
        "-J",
        "--joined_chords",
        metavar="folder",
        nargs="?",
        const="../chords",
        help="Like -C except that all Chords are substituted with the actual Notes they contain. This is useful, "
        "for example, "
        "for relating slurs to the notes they group, or bass figures to their bass notes. ",
    )
    extract_args.add_argument(
        "-D",
        "--metadata",
        metavar="suffix",
        nargs="?",
        const="",
        help="Set -D to update the 'metadata.tsv' files of the respective corpora with the parsed scores. "
        "Add a suffix if you want to update 'metadata{suffix}.tsv' instead.",
    )
    extract_args.add_argument(
        "-p",
        "--positioning",
        action="store_true",
        help="When extracting labels, include manually shifted position coordinates in order to restore them when "
        "re-inserting.",
    )
    extract_args.add_argument(
        "--raw",
        action="store_false",
        help="When extracting labels, leave chord symbols encoded instead of turning them into a single column of "
        "strings.",
    )
    extract_args.add_argument(
        "-u",
        "--unfold",
        action="store_true",
        help="Unfold the repeats for all stored DataFrames.",
    )
    extract_args.add_argument(
        "--interval_index",
        action="store_true",
        help="Prepend a column with [start, end) intervals to the TSV files.",
    )
    extract_args.add_argument(
        "--corpuswise",
        action="store_true",
        help="Parse one corpus after the other rather than all at once.",
    )

    select_facet_args = argparse.ArgumentParser(add_help=False)
    select_facet_args.add_argument(
        "--ask",
        action="store_true",
        help="If several files are available for the selected facet (default: 'expanded', see --use), I will pick one "
        "automatically. Add --ask if you want me to have you select which ones to compare with the scores.",
    )
    select_facet_args.add_argument(
        "--use",
        default="expanded",
        choices=["expanded", "labels"],
        help="Which type of labels you want to compare with the ones in the score. Defaults to 'expanded', "
        "i.e., DCML labels. Set --use labels to use other labels available as TSV and set --ask if several "
        "sets of labels are available that you want to choose from.",
    )

    # main argument parser
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""\
--------------------------
| Welcome to ms3 parsing |
--------------------------

The library offers you the following commands. Add the flag -h to one of them to learn about its parameters.
""",
    )
    parser.add_argument("--version", action="version", version=__version__)
    subparsers = parser.add_subparsers(
        help="The action that you want to perform.", dest="action"
    )

    add_parser = subparsers.add_parser(
        "add",
        help="Add labels from annotation tables to scores.",
        parents=[select_facet_args, parse_args],
    )
    # TODO: staff parameter needs to accept one to several integers including negative
    # add_parser.add_argument('-s', '--staff',
    #                            help="Remove labels from selected staves only. 1=upper staff; -1=lowest staff (
    #                            default)")
    # add_parser.add_argument('--type', default=1,
    #                            help="Only remove particular types of harmony labels.")
    add_parser.add_argument(
        "-s",
        "--suffix",
        metavar="SUFFIX",
        default="_annotated",
        help="Suffix of the new scores with inserted labels. Defaults to _annotated.",
    )
    add_parser.add_argument(
        "--replace",
        action="store_true",
        help="Remove existing labels from the scores prior to adding. Like calling ms3 empty first.",
    )
    add_parser.set_defaults(func=add_cmd)

    check_parser = subparsers.add_parser(
        "check",
        help="""Parse MSCX files and look for errors.
In particular, check DCML harmony labels for syntactic correctness.""",
        parents=[check_args, parse_args],
    )
    check_parser.set_defaults(func=check_cmd)

    compare_parser = subparsers.add_parser(
        "compare",
        help="For MSCX files for which annotation tables exist, create another MSCX file with a coloured label "
        "comparison if differences are found.",
        parents=[select_facet_args, compare_args, parse_args],
    )
    compare_parser.add_argument(
        "-c",
        "--compare",
        metavar="GIT_REVISION",
        default="",
        help="By default, the _reviewed file displays removed labels in red and added labels in green, compared to "
        "the version currently "
        "represented in the present TSV files, if any. If instead you want a comparison with the TSV files from "
        "another Git commit, "
        "pass its specifier, e.g. 'HEAD~3', <branch-name>, <commit SHA> etc.",
    )
    compare_parser.add_argument(
        "-s",
        "--suffix",
        metavar="SUFFIX",
        default="_compared",
        help="Suffix of the newly created comparison files. Defaults to _compared",
    )
    compare_parser.set_defaults(func=compare_cmd)

    convert_parser = subparsers.add_parser(
        "convert",
        help="Use your local install of MuseScore to convert MuseScore files.",
        parents=[parse_args],
    )
    convert_parser.add_argument(
        "-s",
        "--suffix",
        metavar="SUFFIX",
        default="",
        help="Suffix of the converted files. Defaults to " ".",
    )
    convert_parser.add_argument(
        "--format",
        default="mscx",
        help="Output format of converted files. Defaults to mscx. Other options are "
        "{png, svg, pdf, mscz, wav, mp3, flac, ogg, xml, mxl, mid}",
    )
    convert_parser.add_argument(
        "--extensions",
        nargs="+",
        default=["mscx", "mscz"],
        help="Those file extensions that you want to be converted, separated by spaces. Defaults to mscx mscz",
    )
    convert_parser.add_argument(
        "--safe", action="store_false", help="Don't overwrite existing files."
    )
    convert_parser.set_defaults(func=convert_cmd)

    empty_parser = subparsers.add_parser(
        "empty",
        help="Remove harmony annotations and store the MuseScore files without them.",
        parents=[parse_args],
    )
    # TODO: staff parameter needs to accept one to several integers including negative
    # empty_parser.add_argument('-s', '--staff',
    #                            help="Remove labels from selected staves only. 1=upper staff; -1=lowest staff (
    #                            default)")
    # empty_parser.add_argument('--type', default=1,
    #                            help="Only remove particular types of harmony labels.")
    empty_parser.add_argument(
        "-s",
        "--suffix",
        metavar="SUFFIX",
        default="_clean",
        help="Suffix of the new scores with removed labels. Defaults to _clean.",
    )
    empty_parser.set_defaults(func=empty)

    extract_parser = subparsers.add_parser(
        "extract",
        help="Extract selected information from MuseScore files and store it in TSV files.",
        parents=[extract_args, parse_args],
    )
    extract_parser.set_defaults(func=extract_cmd)

    metadata_parser = subparsers.add_parser(
        "metadata",
        help="Update MSCX files with changes made to metadata.tsv (created via ms3 extract -D [-a]).",
        parents=[parse_args],
    )
    metadata_parser.add_argument(
        "-s",
        "--suffix",
        metavar="SUFFIX",
        help="Suffix of the new scores with updated metadata fields.",
    )
    metadata_parser.add_argument(
        "-p",
        "--prelims",
        action="store_true",
        help="Pass this flag if, in addition to updating metadata fields, you also want score headers to be updated "
        "from the columns title_text, subtitle_text, composer_text, lyricist_text, part_name_text.",
    )
    metadata_parser.add_argument(
        "--instrumentation",
        action="store_true",
        help="Pass this flag to update the score's instrumentation based on changed values from "
        "'staff_<i>_instrument' columns.",
    )
    metadata_parser.add_argument(
        "--empty",
        action="store_true",
        help="Set this flag to also allow empty values to be used for overwriting existing ones.",
    )
    metadata_parser.add_argument(
        "--remove",
        action="store_true",
        help="Set this flag to remove non-default metadata fields that are not columns in the metadata.tsv file "
        "anymore.",
    )
    metadata_parser.set_defaults(func=metadata)
    #
    # repair_parser = subparsers.add_parser('repair',
    #                                       help="Apply automatic repairs to your uncompressed MuseScore files.",
    #                                       parents=[parse_args])
    # repair_parser.set_defaults(func=repair)

    review_parser = subparsers.add_parser(
        "review",
        help="Extract facets, check labels, and create _reviewed files.",
        parents=[check_args, select_facet_args, compare_args, extract_args, parse_args],
    )
    review_parser.add_argument(
        "-c",
        "--compare",
        nargs="?",
        metavar="GIT_REVISION",
        const="",
        help="Pass -c if you want the _reviewed file to display removed labels in red and added labels in green, "
        "compared to the version currently "
        "represented in the present TSV files, if any. If instead you want a comparison with the TSV files from "
        "another Git commit, additionally "
        "pass its specifier, e.g. 'HEAD~3', <branch-name>, <commit SHA> etc.",
    )
    review_parser.add_argument(
        "--threshold",
        default=0.6,
        type=float,
        help="Harmony segments where the ratio of non-chord tones vs. chord tones lies above this threshold "
        "will be printed in a warning and will cause the check to fail if the --fail flag is set. Defaults to "
        "0.6 (3:2).",
    )
    review_parser.set_defaults(func=review_cmd)

    transform_parser = subparsers.add_parser(
        "transform",
        help="Concatenate and transform TSV data from one or several corpora. "
        "Available transformations are unfolding repeats and adding an interval index.",
        parents=[parse_args],
    )
    transform_parser.add_argument(
        "-M",
        "--measures",
        action="store_true",
        help="Concatenate measures TSVs for all selected pieces.",
    )
    transform_parser.add_argument(
        "-N",
        "--notes",
        action="store_true",
        help="Concatenate notes TSVs for all selected pieces.",
    )
    transform_parser.add_argument(
        "-R",
        "--rests",
        action="store_true",
        help="Concatenate rests TSVs for all selected pieces (use ms3 extract -R to create those).",
    )
    transform_parser.add_argument(
        "-L",
        "--labels",
        action="store_true",
        help="Concatenate raw harmony label TSVs for all selected pieces (use ms3 extract -L to create those).",
    )
    transform_parser.add_argument(
        "-X",
        "--expanded",
        action="store_true",
        help="Concatenate expanded DCML label TSVs for all selected pieces.",
    )
    transform_parser.add_argument(
        "-F",
        "--form_labels",
        metavar="folder",
        nargs="?",
        const="../form_labels",
        help="Concatenate form label TSVs for all selected pieces.",
    )
    transform_parser.add_argument(
        "-E",
        "--events",
        action="store_true",
        help="Concatenate events TSVs (notes, rests, articulation, etc.) for all selected pieces (use ms3 extract -E "
        "to create those).",
    )
    transform_parser.add_argument(
        "-C",
        "--chords",
        action="store_true",
        help="Concatenate chords TSVs (<chord> tags group notes in the same voice with identical onset and duration) "
        "including "
        "lyrics, dynamics, articulation, staff- and system texts, tempo marking, spanners, and thoroughbass "
        "figures, "
        "for all selected pieces (use ms3 extract -C to create those).",
    )
    transform_parser.add_argument(
        "-D",
        "--metadata",
        action="store_true",
        help="Output 'concatenated_metadata.tsv' with one row per selected piece.",
    )
    transform_parser.add_argument(
        "-s",
        "--suffix",
        nargs="*",
        metavar="SUFFIX",
        help="Pass -s to use standard suffixes or -s SUFFIX to choose your own. In the latter case they will be "
        "assigned to the extracted aspects in the order "
        "in which they are listed above (capital letter arguments).",
    )
    transform_parser.add_argument(
        "-u",
        "--unfold",
        action="store_true",
        help="Unfold the repeats for all concatenated DataFrames.",
    )
    transform_parser.add_argument(
        "--interval_index",
        action="store_true",
        help="Prepend a column with [start, end) intervals to the TSV files.",
    )
    transform_parser.add_argument(
        "--resources",
        action="store_true",
        help="Store the concatenated DataFrames as TSV files with resource descriptors rather than in a ZIP with a "
        "package descriptor.",
    )
    transform_parser.add_argument(
        "--safe", action="store_false", help="Don't overwrite existing files."
    )
    transform_parser.add_argument(
        "--uncompressed",
        action="store_true",
        help="Store the transformed files as uncompressed TSVs rather than writing them into a ZIP file.",
    )
    transform_parser.set_defaults(func=transform_cmd)

    update_parser = subparsers.add_parser(
        "update",
        help="Convert MSCX files to the latest MuseScore version and move all chord annotations "
        "to the Roman Numeral Analysis layer. This command overwrites existing files!!!",
        parents=[parse_args],
    )
    # update_parser.add_argument('-a', '--annotations', metavar='PATH', default='../harmonies',
    #                             help='Path relative to the score file(s) where to look for existing annotation
    #                             tables.')
    update_parser.add_argument(
        "-s",
        "--suffix",
        metavar="SUFFIX",
        help="Add this suffix to the filename of every new file.",
    )
    update_parser.add_argument(
        "--above", action="store_true", help="Display Roman Numerals above the system."
    )
    update_parser.add_argument(
        "--safe",
        action="store_true",
        help="Only moves labels if their temporal positions stay intact.",
    )
    update_parser.add_argument(
        "--staff",
        default=-1,
        help="Which staff you want to move the annotations to. 1=upper staff; -1=lowest staff (default)",
    )
    update_parser.add_argument(
        "--type",
        default=1,
        help="defaults to 1, i.e. moves labels to Roman Numeral layer. Other types have not been tested!",
    )
    update_parser.set_defaults(func=update_cmd)

    return parser

def run():
    parser = get_arg_parser()
    args = parser.parse_args()
    if "func" not in args:
        parser.print_help()
        return
    if args.files is not None:
        args.files = [resolve_dir(path) for path in args.file]
    args.func(args)

if __name__ == "__main__":
    run()


import numpy as np

from PIL import Image

pic = np.load('kittidata/test.npy')
picsouce  = np.load('kittidata/soureces_test.npy')

# for i in range( len(pic)):
# #     ima = Image.fromarray(pic[i],'RGB')
# #     ima.save('kittidata/pic/'+str(i)+'.png')
# #     # ima.show()
this  = picsouce[0]
x = []
len_flames=11
for i in range(len(pic)-len_flames):
    if picsouce[i]==picsouce[i+len_flames]:
        this_flames = pic[i:i+len_flames]
        x.append(this_flames)

x = np.array(x)

np.save('kittidata/test_flames.npy',x)

    # if picsouce[i] !=this:
    #     print(i)
    #     this = picsouce[i]


#!/usr/bin/env python3

from collections import defaultdict
import sys

start = [int(x) for x in sys.stdin.read().strip().split(',')]

start_cnt = defaultdict(int)
rest_cnt = defaultdict(int)

for x in start:
    start_cnt[x] += 1
for _ in range(256):
    #print('start',sorted(start_cnt.items()))
    #print('rest ',sorted(rest_cnt.items()))
    start_cnt_new = defaultdict(int)
    rest_cnt_new = defaultdict(int)
    for age, cnt in start_cnt.items():
        if age == 0:
            start_cnt_new[6] += cnt
            rest_cnt_new[8] += cnt
        else:
            start_cnt_new[age-1] += cnt
    for age, cnt in rest_cnt.items():
        if age == 0:
            start_cnt_new[6] += cnt
            rest_cnt_new[8] += cnt
        else:
            rest_cnt_new[age-1] += cnt
    rest_cnt = rest_cnt_new
    start_cnt = start_cnt_new
a = sum(x for x in start_cnt.values())
b = sum(x for x in rest_cnt.values())
sys.stdout.write(f'{a+b}\n')


import os
import numpy as np
import copy
import torch
import pandas as pd
from torch.utils.data import Subset, Dataset
from torchvision import transforms
from torchvision.datasets import ImageFolder, DatasetFolder
from omegaconf import DictConfig, OmegaConf
from typing import Union, List, Optional, Tuple
import matplotlib.pyplot as plt
from PIL import Image
import logging
from tqdm import tqdm

from .meta.office_caltech import DefaultImageLoader

logger = logging.getLogger(__name__)

from fade.file import FileManager
from . import read_fed_dataset

def simple_stat(arr):
    res = {
        "mean": np.mean(arr),
        "std": np.std(arr),
        "max": np.max(arr),
        "min": np.min(arr),
        "median": np.median(arr),
    }
    for metric, value in res.items():
        print(f"  {metric}: {value:.4g}", end=", ")
    print("")
    return res

def basic_stat_fed_dict(data_dict, verbose=False, group_by="none"):
    users = data_dict['users']
    num_users = len(users)
    num_samples = data_dict['num_samples']
    # if group_by == "group":
    #     groups = data_dict[]

    if verbose:
        example_user = users[0]
        shape = None

        classes = []
        n_classes = []
        mtl_flag = True
        for k in tqdm(users, desc="stat users"):
            if 'idx' in data_dict['user_data'][k]:
                cs = np.unique(data_dict['dataset'].targets[i] for i in data_dict['user_data'][k]['idx']).tolist()
                sample, label = data_dict['dataset'][data_dict['user_data'][k]['idx'][0]]
            elif 'dataset' in data_dict['user_data'][k]:
                # labels = [y for x, y in data_dict['user_data'][k]['dataset']]
                if isinstance(data_dict['user_data'][k]['dataset'], Subset):
                    ds = data_dict['user_data'][k]['dataset'].dataset
                    labels = [ds.targets[i] for i in data_dict['user_data'][k]['dataset'].indices]
                else:
                    ds = data_dict['user_data'][k]['dataset']
                    labels = ds.targets
                labels = np.array(labels)
                cs = np.unique(labels).tolist()
                sample, label = data_dict['user_data'][k]['dataset'][0]

                if len(labels.shape) > 1 and mtl_flag:
                    print(f" Multi-task w/ {labels.shape[1]} tasks")
                    mtl_flag = False
            elif 'x' in data_dict['user_data'][k] and 'y' in data_dict['user_data'][k]:
                cs = np.unique(data_dict['user_data'][k]['y']).tolist()
                sample = data_dict['user_data'][k]['x'][0]
            else:
                raise ValueError(f"Expected keys are: x, y, dataset, idx. Not found "
                                 f"wanted data keys in {data_dict['user_data'][k].keys()}")
            classes.append(cs)
            n_classes.append(len(cs))
            if shape is None:
                shape = sample.shape
            else:
                if shape != sample.shape:
                    logger.warning(f"User {k} has sample of shape not equal to other users' sample "
                                   f"shape: {shape} != {sample.shape} (others)")

            # if len(data_dict['hierarchies']) > 0:
            #     groups = data_dict['hierarchies']
            #     # TODO class balance for each group
            #
            #     # group balance check
            #     uni_groups, n_user_by_group = np.unique(groups, return_counts=True)
            #     print(f" Num user by group: "
            #           f"{dict((g, nu) for g, nu in zip(uni_groups, n_user_by_group))}")
        # classes = np.array(classes)

        print(f" Num of users: {num_users}")
        print(f" Sample: shape {shape}, type {type(sample)}")
        print(f" Num of samples (by user):", end=" ")
        for stat_method in ["sum", "mean", "median", "min", "max", "std"]:
            print(f"{stat_method} {eval('np.'+stat_method+'(num_samples)'):7.1f}", end=", ")
        print()
        if False:  # FIXME: ad-hoc 'float' in str(classes[0].dtype):
            print("mean of objective (per user)", [np.mean(c) for c in classes])
        else:
            classes_ = []
            for cs in classes:
                classes_.extend(cs)
            print(f" Total num of classes: {len(np.unique(classes_))}")
            print(f" Num of classes per user: "
                  f"mean {np.mean(n_classes):.3f}, min {np.min(n_classes)}, max {np.max(n_classes)}")
    return num_samples, num_users

def basic_stat(data_dict, verbose=False):
    """
    Args:
        data_dict: Dict<user_id: Dict<'x': samples, 'y': labels> >.
    """
    num_users = len(data_dict)
    sample = data_dict[list(data_dict.keys())[0]]['x']
    num_samples = np.array([len(data_dict[k]['y']) for k in data_dict])

    if verbose:
        shape = sample.shape[1:]
        classes = [np.unique(data_dict[k]['y']) for k in data_dict]
        n_classes = np.array([len(np.unique(data_dict[k]['y'])) for k in data_dict])

        print(f" Num of users: {num_users}")
        print(f" Sample: shape {shape}, type {type(sample)}")
        print(f" Num of samples (by user):", end=" ")
        for stat_method in ["mean", "median", "min", "max", "std"]:
            print(f"{stat_method} {eval('np.'+stat_method+'(num_samples)'):7.1f}", end=", ")
        if 'float' in str(classes[0].dtype):
            print("mean of objective (per user)", [np.mean(c) for c in classes])
        else:
            print(f" Total num of classes: {len(np.unique(np.concatenate(classes, axis=0)))}")
            print(f" Num of classes per user: "
                  f"mean {np.mean(n_classes):.1f}, min {np.min(n_classes)}, max {np.max(n_classes)}")
    print()
    return num_samples, num_users

def user_class_shards(data_dict, visualize=False, ax=None, show_xticks=True, by_percent=False):
    """Stat the class shards in users.

    Args:
        data_dict: Dict<user_id: Dict<'x': samples, 'y': labels> >.
        visualize: Plot if True.
        by_percent: Transform the count to frequency (percent).

    Returns:
        ucs: An dict<User_name: Dict<class: #samples>> where each element is the number of
            samples of each class in each user.
        ucs_arr: An array of shape (n_user, n_class).
    """
    ucs = dict()
    all_classes = []
    for _id in data_dict:
        y = data_dict[_id]['y']
        classes, cnts = np.unique(y, return_counts=True)
        ucs[_id] = dict((c, cnt) for c, cnt in zip(classes, cnts))
        all_classes = np.union1d(classes, all_classes)
    users = list(data_dict.keys())
    ucs_arr = []
    for _id in users:
        ucs_arr.append([ucs[_id][c] if c in ucs[_id] else 0 for c in all_classes])
    ucs_arr = np.cumsum(np.array(ucs_arr), axis=1)  # shape: (n_user, n_class)
    if by_percent:
        ucs_arr = ucs_arr / ucs_arr[:, -1:] * 100.
    if visualize:
        for ic in np.arange(len(all_classes))[::-1]:
            ax.bar(list(range(len(ucs_arr))), ucs_arr[:, ic], label=f"class-{all_classes[ic]}")
    if show_xticks:
        ax.set(xticks=list(range(len(ucs_arr))))
    return ucs, ucs_arr

def save_data(data_name, data_dict):
    """Save data into files.

    Args:
        data_name ():
        data_dict (dict): Should include keys: 'train' and 'test' at least.
    """
    for s in data_dict:
        fname = FileManager.data(os.path.join(data_name, s, f"{s}.pt"), is_dir=False, overwrite=True)
        with open(fname, 'wb') as outfile:
            print(f"Dumping {s} data => {fname}")
            torch.save(data_dict[s], outfile)

def load_fed_dataset(dataset: str, subset="train",
                     list_keys=("users", "hierarchies", "num_samples"),
                     dict_keys=("user_data",)):
    """Load dataset in fed format without unpacking.
    If multiple files exist, we will load alll and combine into one dict.

    Returns:
        fed_data_dict: A dict like {"users": [], "hierarchies": [], "user_data": {}}
    """
    # os.path.join('data', dataset, 'data', 'train')
    if "-" in dataset:
        dataset, data_fld = dataset.split("-")
    else:
        # data_fld = 'data'
        data_fld = ''

    data_dir = FileManager.data(os.path.join(dataset, data_fld, subset), is_dir=True,
                                      create_new=False)
    # test_data_dir = FileManager.data(os.path.join(dataset, data_fld, 'test'), is_dir=True,
    #                                  create_new=False)
    # clients = []
    # groups = []
    fed_data_dict = {}
    # fed_data_dict = {"users": [], "hierarchies": [], "user_data": {}, "num_samples": []}

    all_files = [f for f in os.listdir(data_dir) if f.endswith('.json') or f.endswith('.pt')]
    for f in all_files:
        file_path = os.path.join(data_dir, f)
        print(f"Load data <== {file_path}")
        if file_path.endswith(".pt"):
            with open(file_path, 'rb') as inf:
                cdata = torch.load(inf)
        elif file_path.endswith(".json"):
            raise ValueError(f"Found json file. Unless you intend to do this, you should use `.pt` file for dataset. "
                             f"`.json` is slow in loading. Try to remove .json file to fix this.")
            # with open(file_path, 'r') as inf:
            #     cdata = json.load(inf)
        else:
            raise TypeError(f"file: {file_path}")
        for k in list_keys:
            if k in cdata:
                fed_data_dict.setdefault(k, []).extend(cdata[k])
        for k in dict_keys:
            if k in cdata:
                fed_data_dict.setdefault(k, {}).update(cdata[k])
        if len(cdata.keys()) > 4:
            raise ValueError(f"Expected three keys. But we get {len(cdata.keys())} "
                             f"keys from the dataset ({dataset}, {subset}): {list(cdata.keys())}")
    return fed_data_dict

def get_dataset_config(name):
    return OmegaConf.load(os.path.join(os.path.dirname(__file__), name + ".yaml")).dataset

def read_data(dataset, load_config=False):
    """parses data in given train and test data directories

    assumes:
    - the data in the input directories are .json files with
        keys 'users' and 'user_data'
    - the set of train set users is the same as the set of test set users

    Args:
        dataset (str|DictConfig):
        load_config (bool): Enforce to use hydra config to access dataset if `dataset`
            is str.

    Return:
        clients: list of client ids
        groups: list of group ids; empty list if none found
        train_data: dictionary of train data
        test_data: dictionary of test data
    """
    logger.warning(f"read_data will be deprecated. Use read_fed_dataset instead.")
    if isinstance(dataset, str) and load_config:
        dataset = get_dataset_config(dataset)
    if isinstance(dataset, DictConfig):
        return read_fed_dataset(dataset)

    train_dict = load_fed_dataset(dataset, "train")
    test_dict = load_fed_dataset(dataset, "test")
    groups = train_dict["hierarchies"] if "hierarchies" in train_dict else []
    if "user_data" in train_dict:
        data_type = "fed"
    # elif "task_data" in train_dict:
    #     data_type = "mtl"
    else:
        raise ValueError(f"train_dict keys: {train_dict.keys()}")
    if data_type == "fed":
        train_data = train_dict["user_data"]

        test_data = test_dict["user_data"]
    # elif data_type == "mtl":
    #     train_data = train_dict["task_data"]
    #
    #     test_data = test_dict["task_data"]
    else:
        raise ValueError(f"data_type: {data_type}")

    users = list(sorted(train_data.keys()))
    # TODO verify the consistency of client keys.

    return users, groups, train_data, test_data

def extract_user_data(user_index, data):  # , dataset):
    """`data` is from the output of `read_fed_data`"""
    if isinstance(data, dict):
        train_data = data["train"]["user_data"]
        test_data = data["test"]["user_data"]
        groups = data["train"]["hierarchies"] if "hierarchies" in data["train"] else []
        users = data["train"]['users']
    else:
        users, groups, train_data, test_data = data
    user_id = users[user_index]
    if len(groups) > 0:
        if isinstance(groups[user_index], list):
            # TODO fix this.
            print(f"WARNING: found vec groups.")
            assert not isinstance(groups[user_index][0], str)
            # FIXME ad-hoc only use the first group if the groups include multiple values.
            group = (groups[user_index][0], groups[user_index][0])
        else:
            group_names = np.unique(groups).tolist()
            group_name = groups[user_index]
            # print(f"### np.unique(groups): {group_names}")
            # note: the unique return SORTED array of groups.
            group = (group_names.index(group_name), group_name)
    else:
        group = (0, 'all')
    data_dict = {'train': train_data, 'test': test_data}
    new_data_dict = {}
    for subset in data_dict:
        subset_data = data_dict[subset]
        if 'idx' in subset_data[user_id]:
            assert isinstance(data, dict), f"data is not dict but {type(data)}"
            assert 'dataset' in data[subset], f"Not found 'dataset' in {subset} data. {data[subset].keys()}"
            dataset = data[subset]['dataset']  # type: Dataset
            idxs = subset_data[user_id]['idx']
            subset_data = Subset(dataset, idxs)
        elif 'dataset' in subset_data[user_id]:
            subset_data = subset_data[user_id]['dataset']
        else:
            subset_data = subset_data[user_id]
            X_set, y_set = subset_data['x'], subset_data['y']
            if not isinstance(X_set, torch.Tensor):
                X_set = torch.Tensor(X_set).type(torch.float32)
                y_set = torch.Tensor(y_set).type(torch.int64)
            subset_data = [(x, y) for x, y in zip(X_set, y_set)]
        new_data_dict[subset] = subset_data
    return user_id, group, new_data_dict['train'], new_data_dict['test']

def grid_imshow(
        tensor: Union[torch.Tensor, List[torch.Tensor]],
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = False,
        img_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: int = 0,
        title=None,
) -> torch.Tensor:
    """Make a grid of images.

    Args:
        tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)
            or a list of images all of the same size.
        nrow (int, optional): Number of images displayed in each row of the grid.
            The final grid size is ``(B / nrow, nrow)``. Default: ``8``.
        padding (int, optional): amount of padding. Default: ``2``.
        normalize (bool, optional): If True, shift the image to the range (0, 1),
            by the min and max values specified by :attr:`range`. Default: ``False``.
        img_range (tuple, optional): tuple (min, max) where min and max are numbers,
            then these numbers are used to normalize the image. By default, min and max
            are computed from the tensor.
        scale_each (bool, optional): If ``True``, scale each image in the batch of
            images separately rather than the (min, max) over all images. Default: ``False``.
        pad_value (float, optional): Value for the padded pixels. Default: ``0``.

    Example:
        See this notebook `here <https://gist.github.com/anonymous/bf16430f7750c023141c562f3e9f2a91>`_

    """
    if not (torch.is_tensor(tensor) or
            (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
        raise TypeError('tensor or list of tensors expected, got {}'.format(type(tensor)))

    # if list of tensors, convert to a 4D mini-batch Tensor
    if isinstance(tensor, list):
        tensor = torch.stack(tensor, dim=0)

    if tensor.dim() == 2:  # single image H x W
        tensor = tensor.unsqueeze(0)
    if tensor.dim() == 3:  # single image
        if tensor.size(0) == 1:  # if single-channel, convert to 3-channel
            tensor = torch.cat((tensor, tensor, tensor), 0)
        tensor = tensor.unsqueeze(0)

    if tensor.dim() == 4 and tensor.size(1) == 1:  # single-channel images
        tensor = torch.cat((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor = tensor.clone()  # avoid modifying tensor in-place
        if img_range is not None:
            assert isinstance(img_range, tuple), \
                "img_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, min, max):
            img.clamp_(min=min, max=max)
            img.add_(-min).div_(max - min + 1e-5)

        def norm_range(t, img_range):
            if img_range is not None:
                norm_ip(t, img_range[0], img_range[1])
            else:
                norm_ip(t, float(t.min()), float(t.max()))

        if scale_each is True:
            for t in tensor:  # loop over mini-batch dimension
                norm_range(t, img_range)
        else:
            norm_range(tensor, img_range)

    if tensor.size(0) == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.size(0)
    xmaps = min(nrow, nmaps)
    ymaps = int(np.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
    num_channels = tensor.size(1)
    fig, axs = plt.subplots(ymaps, xmaps, figsize=(.03 * width * xmaps, .03 * height * ymaps))
    axs = np.atleast_2d(axs)
    # grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            ax = axs[y, x]
            if k < nmaps:
                npimg = tensor[k].numpy()
                ax.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')
                if title is not None:
                    ax.set(title=title(k))
            ax.axis('off')
            # Tensor.copy_() is a valid method but seems to be missing from the stubs
            # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_
            # grid.narrow(1, y * height + padding, height - padding).narrow(  # type: ignore[attr-defined]
            #     2, x * width + padding, width - padding
            # ).copy_(tensor[k])

            k = k + 1
    return fig, axs

class GrayscaleToRgb:
    """Convert a grayscale image to rgb"""
    def __call__(self, image):
        image = np.array(image)
        image = np.dstack([image, image, image])
        return Image.fromarray(image)

class ToBinary:
    """Convert a grayscale image to rgb"""
    def __init__(self, scale=255):
        self.scale = scale

    def __call__(self, image):
        ret = (image > 0).int() * self.scale
        return ret

class SqueezeTarget:
    """Try to squeeze dimensions of targets.
    """

    def __call__(self, y):
        """
        Args:
            y (): Label of a sample.

        Returns:
            Tensor: Converted image.
        """
        return y[0] if len(y) == 1 else y

    def __repr__(self):
        return 'squeeze dimension of targets'

def make_static_dataset_copy(dataset):
    dataset = copy.deepcopy(dataset)
    if isinstance(dataset, Subset):
        ds = dataset.dataset
    else:
        ds = dataset
    if hasattr(ds, 'transform'):
        ds.transform = remove_random_transforms(ds.transform)
    else:
        logger.debug(f"Not found transform. Skip")
    # assert not hasattr(ds, 'target_transform') or ds.target_transform is None
    return dataset

def remove_random_transforms(trans):
    """Remove the random transform in `trans`."""
    if trans is None:
        return trans
    assert isinstance(trans, transforms.Compose)
    new_trans = []
    logger.debug(f"old trans: {trans}")
    for t in trans.transforms:
        if isinstance(t, (transforms.Resize, transforms.ToTensor, transforms.Normalize,
                          DefaultImageLoader)):
            new_trans.append(t)
        elif isinstance(t, transforms.RandomCrop):
            t_ = transforms.CenterCrop(t.size)
            new_trans.append(t_)
            logger.debug(f" Replace {t} by {t_}")
        elif isinstance(t, transforms.RandomResizedCrop):
            t_ = transforms.Resize(t.size)
            new_trans.append(t_)
            logger.debug(f" Replace {t} by {t_}")
        elif isinstance(t, (transforms.RandomHorizontalFlip, transforms.RandomVerticalFlip, transforms.RandomRotation)):
            logger.debug(f" Skip transform: {t}")
        else:
            raise TypeError(f"Unable to translate transform: {t}")
    new_trans = transforms.Compose(new_trans)
    logger.debug(f"New trans: {new_trans}")
    return new_trans

def update_dataset_targets(dataset, new_targets, indices=None):
    """Update dataset targets in place."""
    if indices is None:
        indices = list(range(len(dataset)))
    if isinstance(dataset, Subset):
        indices = dataset.indices
        update_dataset_targets(dataset.dataset, new_targets, indices=indices)
        # for ii in indices:
        #     dataset.dataset.targets[indices[ii]] = new_targets[ii]
    elif isinstance(dataset, DatasetFolder):
        assert len(indices) == len(new_targets), f"len(indices) ({len(indices)}) != len(new_targets) ({len(new_targets)})"
        for i in range(len(indices)):
            s = dataset.samples[indices[i]]
            dataset.samples[indices[i]] = (s[0], new_targets[i])
        # for i in range(len(dataset.samples)):
        #     dataset.samples[indices[i]] = (s[0], 0)
        dataset.targets = [s[1] for s in dataset.samples]
    elif isinstance(dataset, list) and isinstance(dataset[0], tuple) and len(dataset[0])==2:
        for i in range(len(indices)):
            s = dataset[indices[i]]
            dataset[indices[i]] = (s[0], new_targets[i])
    else:
        raise TypeError(f"Type: {type(dataset)}")
        # TODO use this for other known types. Do not use this for all Dataset.
        # assert hasattr(dataset, 'targets')
        # for ii in indices:
        #     dataset.targets[indices[ii]] = new_targets[ii]

def merge_fed_dataset_users(dataset):
    """Merge all users into a single table."""
    assert isinstance(dataset, dict)

    dfs = dict((k, []) for k in ['x', 'y', 'group', 'user'])  # a list of all users' dataframes
    for user, group in zip(dataset['users'], dataset['hierarchies']):
        X = dataset['user_data'][user]['x']
        y = dataset['user_data'][user]['y']
        df = {'x': X.tolist(), 'y': y.tolist(), 'group': [group]*len(y), 'user': [user]*len(y)}
        for k in df:
            dfs[k].extend(df[k])
    return dfs

def compute_mean_distance(data_dict):
    users = data_dict["users"]
    n_user = len(users)
    D = np.zeros((n_user, n_user))
    for i in range(n_user):
        X1 = data_dict['user_data'][users[i]]['x']
        for j in range(i+1, n_user):
            X2 = data_dict['user_data'][users[j]]['x']
            d = np.linalg.norm(np.mean(X1, axis=0) - np.mean(X2, axis=0))
            D[i, j] = d
            D[j, i] = d
    return D

def plot_sample_size_dist(fed_dict, subset):
    users = fed_dict[subset]["users"]
    n_user = len(fed_dict[subset]["users"])
    from .utils import extract_user_data
    n_smp_by_class = {}
    for i_user in range(n_user):
        user_id, group, train_set, test_set = extract_user_data(i_user, fed_dict)
        subset_data = train_set if subset == "train" else test_set
        for _, y in subset_data:
            n_smp_by_class.setdefault(y, [0] * n_user)
            n_smp_by_class[y][i_user] += 1
    all_classes = sorted(list(n_smp_by_class.keys()))
    n_smp_mat = np.array([n_smp_by_class[cl] for cl in all_classes]).T
    # print(n_smp_mat)
    m_classes, m_users = np.meshgrid(all_classes, np.arange(0, n_user))
    m_classes = m_classes.reshape((-1,))
    m_users = m_users.reshape((-1,))
    n_smp_mat = n_smp_mat.reshape((-1,))

    import seaborn as sns
    import matplotlib.pyplot as plt
    max_size = 240
    n_smp_mat = n_smp_mat * 1. * max_size / np.max(n_smp_mat)
    fig, ax = plt.subplots(1, 1)
    plt.scatter(m_classes, m_users, s=n_smp_mat, alpha=0.5)
    # sns.scatterplot(x=m_classes, y=m_users, size=n_smp_mat, ax=ax)
    ax.set(xlabel="class", ylabel="user", yticks=list(range(n_user)), yticklabels=users)
    return ax

def get_dataloader_options(dataset):
    if isinstance(dataset, Subset):
        dataset = dataset.dataset
    return dataset.dataloader_options() \
        if hasattr(dataset, "dataloader_options") else {}


'''
Write a script that takes in a list of numbers and:
    - sorts the numbers
    - stores the numbers in tuples of two in a list
    - prints each tuple

If the user enters an odd numbered list, add the last item
to a tuple with the number 0.

Note: This lab might be challenging! Make sure to discuss it with your mentor
or chat about it on our forum.

'''

samplist = [2, 5, 7, 98, 1, 45, 9, 3]
samplist.sort()
biglist = []
ending = samplist[-1]

if len(samplist) % 2:
    x = 0
    while x < len(samplist):
        biglist.append(tuple(samplist[ x, x + 1]))
        x += 2

else:
    x = 0
    while x < len(samplist):
        biglist.append(tuple(samplist[x, x + 1]))
        x += 2

print(biglist)


#함수형모델 앙상블 + RMSE, R2넣기
# 데이터가 x 는 2개, y는 1개인 모델 output은 1개
#1. 데이터

import numpy as np
#모델1.
x1 = np.array([range(1,101), range(711,811), range(100)])
x2 = np.array([range(4,104), range(761,861), range(100)])

#모델2.
y1 = np.array([range(101,201), range(311, 411), range(100)])

x1 = np.transpose(x1)
x2 = np.transpose(x2)

y1 = np.transpose(y1)

#train, test분리 (슬라이싱 대신에 train_test_split)
from sklearn.model_selection import train_test_split
x1_train, x1_test, y1_train, y1_test = train_test_split(x1, y1, shuffle=True, train_size =0.7)
x2_train, x2_test = train_test_split(x2,shuffle=True, train_size =0.7)

print("x1_test", x1_test.shape)
print("x2_test", x2_test.shape)
print("y1_test", y1_test.shape)

#함수형 모델 만들기
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Input

#모델1.
input1 = Input(shape =(3,))
dense1 = Dense(10, activation='relu', name='king1')(input1) #모델 이름이 중복될 때 , 혹은 dense이름 변경
dense2 = Dense(7, activation='relu', name='king2')(dense1)
dense3 = Dense(5, activation='relu', name='king3')(dense2)
output1 = Dense(3, activation ='linear', name='king4')(dense3)

#모델2.
input2 = Input(shape =(3,))
dense1_1 = Dense(15, activation='relu')(input2)
dense2_1 = Dense(11, activation='relu')(dense1_1)
output2 = Dense(3, activation ='linear')(dense2_1)

#모델끼리는 별 영향을 끼치지 않는다
##########모델 병합, concantenate 여러가지 방법으로import
from tensorflow.keras.layers import Concatenate, concatenate
# from keras.layers.merge import Concatenate, concatenate
# from keras.layers import Concatentate, concatenate

#########모델 2개 병합하기
#소문자 concatenate
#merge1 = concatenate([output1, output2])#리스트로 묶기

#방법 1. Concatenate대문자 사용법
#merge1 = Concatenate()([output1, output2]) #대문자 Concatenate 면 class다
#방법 2.
merge1 = Concatenate(axis=1)([output1, output2]) #axis축 // 여기서의 output은 y축의 데이터 개수가 아니고 모델 1의output, 모델2의output
middle1 = Dense(30)(merge1) #함수형으로 모델 병합하기
middle1 = Dense(7)(middle1)
middle1 = Dense(11)(middle1)

'''
merge1 = concatenate([output1, output2])#리스트로 묶기
#앞에 변수명 같이 명시 해도 가능.
middle1 = Dense(30)(merge1) #함수형으로 모델 병합하기
middle2 = Dense(7)(middle1)
middle3 = Dense(11)(middle2) #이렇게 middle1도 가능?
'''
########### output 만들기 (분기) y갯수에 따라서
output3 = Dense(30)(middle1)#middle1에서 받아온다
output3_1 = Dense(7)(output3)
output3_2 = Dense(3)(output3_1)#outputs에 마지막 아웃풋 넣기

# output4 = Dense(15)(middle1)
# output4_1 = Dense(14)(output4)
# output4_2 = Dense(11)(output4_1)
# output4_3 = Dense(3)(output4_2)#마지막 아웃풋 넣기

#모델 정의하기
model = Model (inputs =[input1, input2], outputs =output3_2)
                        #모델 1,2넣기               #ㄴ여기아웃풋은 input할 때의 output이 아니라 output할 때 마지막 output변수명
model.summary()

#concatenate는 연산하지 않고 병합만 해준다

# 컴파일, 훈련
model.compile(loss= 'mse', optimizer='adam', metrics=['mae'])
model.fit([x1_train, x2_train], y1_train, epochs=100, batch_size=8, validation_split=0.25, verbose=1)

# 평가
result = model.evaluate([x1_test, x2_test], y1_test, batch_size=8)
# (loss, 메트릭스의 mse)가 나온다

print('result:' , result)
#result: [1.739203691482544, 0.4144916236400604, 1.3247121572494507, 0.47640863060951233, 0.80832839012146]

#예측 = 훈련시킨 값이 나온다 (accuracy 의 정확도를 위해서 예측이 필요하다) x값을 넣어서 예측 y값을 나오게
y_predict = model.predict([x1_test, x2_test]) #y_predict라는 변수에 위에서 validation_split으로 쪼갠 test데이터에서 y값을 예측하자
print('결과물: \n :' , y_predict)

#RMSE
#사이킷런 지표
from sklearn.metrics import mean_squared_error
def RMSE(y1_test, y_predict) : #y1_test validation_split값의 y예측 테스트 데이터를 가져오고, 예측 데이터를 넣을 변수명
    return np.sqrt(mean_squared_error(y1_test, y_predict))
print('RMSE:' ,RMSE(y1_test, y_predict))

#R2
from sklearn.metrics import r2_score
r2 = r2_score(y1_test, y_predict)
print('R2:', r2)


# general imports
import os
import sys
from omegaconf import DictConfig, OmegaConf
import hydra

# smartsim and smartredis imports
from smartsim import Experiment
from smartsim.settings import MpiexecSettings

## Define function to parse node list
def parseNodeList(fname):
    with open(fname) as file:
        nodelist = file.readlines()
        nodelist = [line.rstrip() for line in nodelist]
        nodelist = [line.split('.')[0] for line in nodelist]
    nNodes = len(nodelist)
    return nodelist, nNodes

## Co-located DB launch
def launch_coDB(cfg, nodelist, nNodes):
    # Print nodelist
    print(f"\nRunning on {nNodes} total nodes on Polaris")
    print(nodelist, "\n")
    hosts = ','.join(nodelist)

    # Initialize the SmartSim Experiment
    PORT = cfg.experiment.port
    exp = Experiment(cfg.experiment.name, launcher=cfg.experiment.launcher)

    # Set the run settings, including the Python executable and how to run it
    Py_exe = cfg.experiment.sim_executable
    exe_args = Py_exe + f' --dbnodes 1' \
                      + f' --ppn {cfg.run_args.simprocs_pn} --logging {cfg.logging}'
    run_settings = MpiexecSettings(
                   'python',
                   exe_args=exe_args,
                   run_args=None,
                   env_vars=None
                   )
    run_settings.set_tasks(cfg.run_args.simprocs)
    run_settings.set_tasks_per_node(cfg.run_args.simprocs_pn)
    run_settings.set_hostlist(hosts)
    run_settings.set_cpu_binding_type(cfg.run_args.sim_cpu_bind)

    # Create the co-located database model
    colo_model = exp.create_model("load_data", run_settings)
    #kwargs = {
    #    maxclients: 100000,
    #    threads_per_queue: 1, # set to 4 for improved performance
    #    inter_op_threads: 1,
    #    intra_op_threads: 1,
    #    #server_threads: 2 # keydb only
    #    }
    colo_model.colocate_db(
            port=PORT,
            db_cpus=cfg.run_args.dbprocs_pn,
            debug=True,
            limit_app_cpus=False,
            ifname=cfg.run_args.network_interface,
            )

    # Start the co-located model
    exp.start(colo_model, block=False, summary=True)

    # Set up the data consumer (i.e. training)
    ml_exe = cfg.experiment.ml_executable
    exe_args = ml_exe + f'  --dbnodes 1 --device {cfg.run_args.device}' \
               + f' --ppn {cfg.run_args.mlprocs_pn} --logging {cfg.logging}'
    run_settings_ML = MpiexecSettings(
            'python',
            exe_args=exe_args,
            run_args=None,
            env_vars=None
            )
    print(cfg.run_args.mlprocs,cfg.run_args.mlprocs_pn,hosts,cfg.run_args.ml_cpu_bind)
    run_settings_ML.set_tasks(cfg.run_args.mlprocs)
    run_settings_ML.set_tasks_per_node(cfg.run_args.mlprocs_pn)
    run_settings_ML.set_hostlist(hosts)
    run_settings_ML.set_cpu_binding_type(cfg.run_args.ml_cpu_bind)

    # Start the data consumer
    ml_model = exp.create_model("train_model", run_settings_ML)
    exp.start(ml_model, block=True, summary=True)

## Main function
@hydra.main(version_base=None, config_path="../conf", config_name="config")
def main(cfg: DictConfig):
    # Get nodes of this allocation (job)
    hostfile = os.getenv('PBS_NODEFILE')
    nodelist, nNodes = parseNodeList(hostfile)

    # Call appropriate launcher
    if (cfg.database.launch == "colocated"):
        launch_coDB(cfg,nodelist,nNodes)
    elif (cfg.database.launch == "clustered"):
        launch_clDB(cfg,nodelist,nNodes)
    else:
        print("\nERROR: Launcher is either colocated or clustered\n")

    # Quit
    print("Done")
    print("Quitting")

## Run main
if __name__ == "__main__":
    main()


from pyparsing import alphas
from pyparsing import alphanums
from pyparsing import LineStart
from pyparsing import LineEnd
from pyparsing import Literal
from pyparsing import nums
from pyparsing import OneOrMore
from pyparsing import Optional
from pyparsing import SkipTo
from pyparsing import StringEnd
from pyparsing import Word
from pyparsing import ZeroOrMore

charge_mapping = {
    '0': 0,
    '1': 3,
    '2': 2,
    '3': 1,
    '4': None,
    '5': -1,
    '6': -2,
    '7': -3
}

atom_stereo_parity_mapping = {
    '0': 'not stereo',
    '1': 'odd',
    '2': 'even',
    '3': 'either or unmarked stereo center'
}

hydrogen_count_mapping = {
    '0': None,
    '1': 'H0',
    '2': 'H1',
    '3': 'H2',
    '4': 'H3',
    '5': 'H5'
}

bond_type_mapping = {
    '1': 'Single',
    '2': 'Double',
    '3': 'Triple',
    '4': 'Aromatic',
    '5': 'Single or Double',
    '6': 'Single or Aromatic',
    '7': 'Double or Aromatic',
    '8': 'Any'
}

bond_stereo_mapping = {
    'Single': {
        '0': 'Not stereo',
        '1': 'Up',
        '4': 'Either',
        '6': 'Down'
    },
    'Double': {
        '0': 'Use x-, y-, z-coords from atom block to determine cis or trans',
        '3': 'Either cis or trans double bond'
    }
}

bond_topology_mapping = {
    '0': 'Either',
    '1': 'Ring',
    '2': 'Chain'
}

reacting_center_status_mapping = {
    '0': 'unmarked',
    '1': 'a center',
    '-1': 'not a center',
    '2': 'no change',
    '4': 'bond made/broken',
    '8': 'bond order changes',
    '12': 'bond made/broken + bond order changes',
    '5': 'bond made/broken + a center',
    '13': 'bond made/broken + bond order changes + a center'
}

class SDFParser(object):

    some_line = LineStart() + SkipTo(LineEnd()).suppress()
    empty_line = LineStart() + LineEnd()

    num_atoms = Word(nums)
    num_bonds = Word(nums)
    num_atom_lists = Word(nums)
    obsolete1 = Word(nums)
    molecule_is_chiral = Word(nums)
    num_stext_entries = Word(nums)
    obsolete2 = Word(nums)
    obsolete3 = Word(nums)
    obsolete4 = Word(nums)
    obsolete5 = Word(nums)
    num_lines_with_add_properties = Word(nums)
    version = Word(alphanums)

    counts_line = num_atoms + num_bonds + num_atom_lists + \
        obsolete1 + molecule_is_chiral + num_stext_entries + obsolete2 + \
        obsolete3 + obsolete4 + obsolete5 + num_lines_with_add_properties + \
        version + LineEnd().suppress()

    atom_x_coordinate = Word(nums + '-.')
    atom_y_coordinate = Word(nums + '-.')
    atom_z_coordinate = Word(nums + '-.')
    # entry in periodic table or L for atom list, A, Q, * for unspecified atom,
    # and LP for lone pair, or R# for Rgroup label
    atom_symbol = Word(alphas)
    # -3, -2, -1, 0, 1, 2, 3, 4 (0 if value beyond these limits)
    atom_mass_difference = Word(nums + '-')
    # 0, 1, 2, 3, 4, 5, 6, 7 (0 if uncharged; 1, 2, 3 if positive charges where
    # 1 = +3, 2 = +2, 3 = +1; 4 if doublet radical; 5, 6, 7 if negative charges
    # where 5 = -1, 6 = -2, 7 = -3)
    # Wider range of values in M CHG and M RAD lines below. Retained for
    # compatibility with older Ctabs, M CHG and M RAD lines take precedence.
    atom_charge = Word(nums + '-')
    # 0 = not stereo, 1 = odd, 2 = even, 3 = either or unmarked stereo center
    atom_stereo_parity = Word(nums)
    # 1 = H0, 2 = H1, 3 = H2, 4 = H3, 5 = H4
    # H0 means no H atoms allowed unless explicitly drawn. Hn means atom must
    # have n or more H’s in excess of explicit H’s.
    hydrogen_count_plus_1 = Word(nums)
    # 0 = ignore stereo configuration of this double bond atom, 1 = stereo
    # configuration of double bond atom must match
    # Double bond stereochemistry is considered during atom_stereo_parity only
    # if both ends of the bond are marked with stereo care boxes.
    stereo_care_box = Word(nums)
    # 0 = no marking (default) (1 to 14) = (1 to 14) 15 = zero valence
    # Shows number of bonds to this atom, including bonds to implied H’s.
    valence = Word(nums)
    # 0 = not specified, 1 = no H atoms allowed
    # Redundant with hydrogen count information. Might be unsupported in future
    # releases of Symyx software.
    h0_designator = Word(nums)
    not_used1 = Word(nums)
    not_used2 = Word(nums)
    # 1 - number of atoms
    atom_atom_mapping_number = Word(nums)
    # 0 = property not applied, 1 = configuration is inverted,
    # 2 = configuration is retained
    inversion_retention_flag = Word(nums)
    # 0 = property not applied, 1 = change on atom must be exactly as shown
    exact_change_flag = Word(nums)

    atom_block_row = atom_x_coordinate + atom_y_coordinate + \
        atom_z_coordinate + atom_symbol + atom_mass_difference + \
        atom_charge + atom_stereo_parity + hydrogen_count_plus_1 + \
        stereo_care_box + valence + h0_designator + not_used1 + not_used2 + \
        atom_atom_mapping_number + inversion_retention_flag + \
        exact_change_flag + LineEnd().suppress()

    atom_block = OneOrMore(atom_block_row)

    # 1 - number of atoms
    first_atom_number = Word(nums)
    # 1 - number of atoms
    second_atom_number = Word(nums)
    # 1 = Single, 2 = Double, 3 = Triple, 4 = Aromatic, 5 = Single or Double,
    # 6 = Single or Aromatic, 7 = Double or Aromatic, 8 = Any
    bond_type = Word(nums)
    # Single bonds: 0 = not stereo, 1 = Up, 4 = Either, 6 = Down, Double bonds:
    # 0 = Use x-, y-, z-coords from atom block to determine cis or trans,
    # 3 = Cis or trans (either) double bond
    # The wedge (pointed) end of the stereo bond is at the first atom (Field
    # first_atom_number above)
    bond_stereo = Word(nums)
    not_used3 = Word(nums)
    # 0 = Either, 1 = Ring, 2 = Chain
    bond_topology = Word(nums)
    # 0 = unmarked, 1 = a center, -1 = not a center, Additional: 2 = no change,
    # 4 = bond made/broken, 8 = bond order changes 12 = 4+8 (both made/broken
    # and changes); 5 = (4 + 1), 9 = (8 + 1), and 13 = (12 + 1) are also
    # possible
    reacting_center_status = Word(nums)
    bound_block_row = first_atom_number + second_atom_number + bond_type + \
        bond_stereo + not_used3 + bond_topology + reacting_center_status + \
        LineEnd().suppress()

    bound_block = ZeroOrMore(bound_block_row)

    # M  CHG  6  18  -1  27  -1  35  -1  36   1  37   1  38   1
    #         ^  `--v--´ `--v--´
    #         |  entry 1 entry 2 ...
    #         6 entries following
    # entry 1: atom number 18 has charge -1
    # entry 2: atom number 27 has charge -1
    # ...
    num_charge_entries = Word(nums)
    charge_atom_number = Word(nums)
    charge_atom_charge = Word(nums + '-')
    charge_props_line = Literal('M  CHG') + num_charge_entries + \
        OneOrMore(charge_atom_number + charge_atom_charge) + \
        LineEnd().suppress()
    properties_block = ZeroOrMore(charge_props_line)
    blocks_end_marker = 'M  END' + LineEnd()

    key = Word(alphanums + '_')
    val = OneOrMore(Word(alphanums + '?&.,:;_=/\\\'-+*#@()[]{}'))
    sdf_item = Literal('>  <').suppress() + key + Literal('>').suppress() + \
        LineEnd().suppress() + val + LineEnd().suppress() + \
        empty_line.suppress()

    compound = counts_line + atom_block + bound_block + \
        Optional(properties_block) + blocks_end_marker.suppress() + \
        OneOrMore(sdf_item)

    compound_delimiter = Literal('$$$$')
    sdf = OneOrMore(compound + compound_delimiter.suppress()) + \
        StringEnd().suppress()

    def __init__(self):
        self.counts_line.addParseAction(self._counts_line_hook)
        self.atom_block.addParseAction(self._atom_block_hook)
        self.bound_block.addParseAction(self._bound_block_hook)
        self.properties_block.addParseAction(self._properties_block_hook)
        self.blocks_end_marker.addParseAction(self._blocks_end_hook)
        self.sdf_item.addParseAction(self._sdf_item_hook)
        self.sdf.addParseAction(self._finalize_sdf_items_hook)
        self._converters = []

    def _counts_line_hook(self, parsing_results):
        for converter in self._converters:
            converter.convert_counts_line(parsing_results)

    def _atom_block_hook(self, parsing_results):
        for converter in self._converters:
            converter.read_atom_block(parsing_results)

    def _bound_block_hook(self, parsing_results):
        for converter in self._converters:
            converter.convert_bound_block(parsing_results)

    def _properties_block_hook(self, parsing_results):
        for converter in self._converters:
            converter.update_atom_block(parsing_results)

    def _blocks_end_hook(self):
        for converter in self._converters:
            converter.finalize_atom_block_conversion()

    def _sdf_item_hook(self, parsing_results):
        for converter in self._converters:
            converter.read_sdf_item(parsing_results)

    def _finalize_sdf_items_hook(self, parsing_results):
        for converter in self._converters:
            converter.finalize_sdf_items()

    def parse(self, file_path):
        self.sdf.parseFile(file_path)

    def register_converter(self, converter):
        """
        :param converter: a subclass of SDFConverter
        """
        self._converters.append(converter)


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

import matplotlib.pyplot as plt
import numpy as np
import glob
import cv2
from PIL import Image
import argparse

from models import *
from utils import *

# All implemented models
models = ['resnet', 'vgg16', 'googlenet', 'resnext', 'SimpleResNeXt_v1', 'SimpleResNeXt_v2']
parser = argparse.ArgumentParser()
parser.add_argument("-model", "--model-name", help="model name", dest="model_name", default='SimpleResNeXt_v1', choices=(tuple(models)))
args = parser.parse_args()

model_name = args.model_name

# Cifar-10 labels
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# Use GPU or not
USE_GPU = True if torch.cuda.is_available() else False

# Preprocessing
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

# Load Model
net = load_model(model_name, USE_GPU)

# Where are the images & Read them
img_folder = './test_imgs/'
imgs = []
for img_path in glob.glob(img_folder+'*.png'):
	img = Image.open(img_path)
	imgs.append(img)

# predict each image
for i, img in enumerate(imgs):
	mat = np.asarray(img)
	x = transform_test(mat)
	x = x.unsqueeze(0)

	if USE_GPU:
		x = x.cuda()
	with torch.no_grad():
		net.eval()
		out = net(x)
		top_5_score, top_5_label = out.topk(5)
		top_5 = [classes[l] for i, l in enumerate(top_5_label.data.cpu().numpy()[0])]
		title = [str(i+1) + ' ' + label + '\n' for i, label in enumerate(top_5)]
		title = ''.join(title)
		plt.subplot(2, 5, i+1), plt.imshow(img), plt.title(title), plt.yticks([]), plt.xticks([])
	if i > 10 :
		break
plt.show()


import sys
import re
import time
import requests

from utils import utils
from bs4 import BeautifulSoup

import modules.psql.political_committees as political_committees_model
import modules.psql.alternate_names as alternate_names_model
# import modules.psql.political_party_units as political_party_units_model

def scrape_page(html, party, year):
    soup = BeautifulSoup(html, "html.parser")
    tables = get_tables(soup)
    table_counter = 0
    for table in tables:
        table_counter+=1
        if table_counter < 4:
            continue
        row_counter = 0
        rows = get_rows(table)
        for row in rows:
            row_counter+=1
            data = get_data(row)
            # if row_counter > 10:
            #     continue
            process_data(data, party, year)

def get_tables(soup):
    return soup.findAll('table')

def get_rows(soup):
    return soup.findAll('tr')

def get_data(soup):
    return soup.findAll('td')

def process_data(data, party, year):
    record_obj = build_record(data, party, year)
    if record_obj and record_obj['name'] != 'Committee':
        print ("year: %s | name: %s " % (year, record_obj['name']))
        result = political_committees_model.touch_political_committee(record_obj)
        if result['slug_name'] != record_obj['slug_name']:
            alternate_names_model.touch_alternate_name(record_obj)

def build_record(data, party, year):
    # 2006 no reports...
    if int(year) > 2006:
        committee, address, telephone, registration_id, registration_date, term_date, reports = [(c.text).strip() for c in data]
    else:
        committee, address, telephone, registration_id, registration_date, term_date = [(c.text).strip() for c in data]
    result = {
        "name": committee,
        "slug_name": utils.slugify_string(committee),
        "registration_id": registration_id,
        "party": party,
        "created_at": "now()",
        "updated_at": "now()"
    }
    return result

def main(args):

    pages_to_scrape = [
        {'year':'2016', 'm': 'apr 04', 'url': 'https://web.archive.org/web/20160404112756/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2015', 'm': 'may 24', 'url': 'https://web.archive.org/web/20150524091118/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2014', 'm': 'may 08', 'url': 'https://web.archive.org/web/20140508022548/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2013', 'm': 'may 03', 'url': 'https://web.archive.org/web/20130503144622/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2012', 'm': 'nov 04', 'url': 'https://web.archive.org/web/20121104205208/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2011', 'm': 'aug 21', 'url': 'https://web.archive.org/web/20110821014942/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2010', 'm': 'may 27', 'url': 'https://web.archive.org/web/20100527154013/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2009', 'm': 'may 24', 'url': 'https://web.archive.org/web/20090524061459/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2008', 'm': 'mar 17', 'url': 'https://web.archive.org/web/20080317133140/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2007', 'm': 'aug 23', 'url': 'https://web.archive.org/web/20070823030557/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2006', 'm': 'apr 27', 'url': 'https://web.archive.org/web/20060427064942/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2005', 'm': 'feb 10', 'url': 'https://web.archive.org/web/20050210035907/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2005', 'm': 'feb 10', 'url': 'https://web.archive.org/web/20050210035907/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''},
        {'year':'2004', 'm': 'feb 10', 'url': 'https://web.archive.org/web/20040801000000*/http://www.cfboard.state.mn.us/campfin/pcfatoz.html', 'party':''}
    ]

    # pages_to_scrape = [pc, ppu]
    for page in pages_to_scrape:
        r = requests.get(page['url'])
        html = r.text
        scrape_page(html, page['party'], page['year'])

if __name__ == '__main__':
    main(sys.argv)

# docker run -it --rm --name mncf-scripts --env-file ./dev/dev.env -v "$PWD":/usr/src/app --link postgres-mncf:postgres python/mncf-scripts python ./scripts/get_political_groups.py


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Nov 25 15:12:34 2019

@author: marco
"""
import random
import matplotlib.pyplot as plt
import time
import numpy as np

def counting_sort_integers(input_int_list, max_int):

    counter = [0] * (max_int + 1)

    # Counting the occurrences of a number in the input list
    for n in input_int_list:
        counter[n] += 1

    sorted_list = [0] * len(input_int_list) # Default value for sorted list is not in range [0-9] (debug purposes)
    current_index = 0 # We have to keep track of the next available index in the sorted list

    for i in range(len(counter)): # We have to use values in range [0-len(counter)]-> [0-9].
        count = counter[i] # How many occurrences have there been in the input list?
        if count > 0: # If there is something to put in the sorted list
            sorted_list[current_index:(current_index+count)] = [i] * count # Then put them (count) at a times
            current_index += count

    return sorted_list

def counting_sort_list_letters(input_str):
    ''' The counting sort here is implemented using the feature that python dictionaries
        provide: ordering of the keys while preserving access in constant time.
    '''

    # Initializing the dict that will act as counter, access will be in constant time.
    letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
    counter_letters = dict.fromkeys(list(letters), 0)

    # Counting the occurrences of the letters in the input string (k)
    for l in list(input_str):
        counter_letters[l] += 1

    # Initializing the value of the sorted array to a default char not in alphabet (debug purposes, could be any char value)
    sorted_array = ['?'] * len(input_str)

    # This is here to keep track of how many chars have been inserted in the string, so we can access using
    # this variable as an index.
    current_index = 0
    for k in counter_letters.keys(): # Keys list keeps the order, now traversing it from left to right (from the least important value to the most important)

        count = counter_letters[k] # First get how many occurrences of a letter have been found

        if count > 0: # If > 0 then the letter (key) will end up in the sorted array (count) times
            sorted_array[current_index:(current_index+count)] = [k] * count # Put multiple separated values at multiple separated indexes (1 index -> 1 value)
            counter_letters[k] -= count # Just keeping track of the inserted values
            current_index += count # The position has to move of (count) moves to the right

    return sorted_array

def counting_sort_chars(list_of_chars):
    mapping = [ord(c) for c in list_of_chars]
    mapping = counting_sort_integers(mapping, max(mapping))
    return [chr(x) for x in mapping]

''' This implementation works, but it's absolutely inefficient resources-wise.
    The keys (integers) used in the counting sort for strings become very large, very quickly.

def map_string_to_int(string): # Mapping to int allows for an ordering of strings.
    res = 0
    for i in range(len(string)):
       index = len(string) - i - 1 # We start mapping from the most significant char (so from the right)
       res += map_char_to_int(string[index]) * (52 ** i)

    return res

def count_sort_list_of_strings_inefficient(strings, length_m, string_len_n):
    len_list = [[]for i in range(string_len_n + 1)]
    res = []

    for s in strings:
            len_list[len(s)][map_string_to_int(s)] = s

    # Since we are using the dictionaries, where the key is the mapping from string to int, we sort the keys.
    for d in len_list:
        if d != {}:
            k = d.keys()
            ordered_k = counting_sort_integers(k)

            for ok in ordered_k:
                res.append(d[ok])

    return res
'''

def radix_sort(strings, i):
    letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
    res = []

    if len(strings) < 2:
        return strings
    else:
        s_dict = {k:[] for k in letters}

        for s in strings:
            s_dict[s[i]].append(s)

        for k in s_dict.keys():
            if s_dict[k] != []:
                res.extend(radix_sort(s_dict[k], i+1))

    return res

def counting_sort_strings(strings, m_num_strings, n_max_string_length):

    subgroups = [[] for _ in range(n_max_string_length + 1)]

    for string in strings:
        subgroups[len(string)-1].append(string)

    res = []
    for s in subgroups:
        if len(s) > 0:
            tmp_r = radix_sort(s, 0)
            res.extend(tmp_r)

    return res

def test_counting_sort_for_integers_chars():
    ''' Empirical results from counting sort for single words'''

    strings = []
    letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'

    # Generate random strings of length from 1 to (iterations)
    iterations = 5000
    for i in range(1, iterations+1):
        curr_str = ''
        for j in range(0, i):
            letter = letters[random.randint(0, len(letters)-1)]
            curr_str += letter

        strings.append(curr_str)

    performance = []
    for i in range(len(strings)):
        s = strings[i]
        start = time.time()
        counting_sort_list_letters(s)
        end = time.time()
        elapsed_time = end - start
        performance.append([elapsed_time, s])

    #performance.sort(key=lambda x : x[0], reverse=True)

    plt.figure(figsize=(10, 7))
    plt.xlabel('String length')
    plt.ylabel('Time elapsed (s)')

    times = [x[0] for x in performance]
    plt.plot(times)
    plt.show()

def test_counting_sort_strings():
    ''' Empirical results from counting sort applied to list of strings '''
    s = ['a', 'ab', 'bab', 'aba','b']
    m = len(s)
    n = 3
    r = counting_sort_strings(s, m, n)

    return r

if __name__ == '__main__':
    res = test_counting_sort_strings()


import os
import pyNN.spiNNaker as sim
from spinnaker_testbase import BaseTestCase
from spynnaker.pyNN.data import SpynnakerDataView

class TestNoIobufDuringRun(BaseTestCase):

    def check_for_iobufs(self, prov_path):
        return any("iobuf" in filename and ".txt" in filename
                   for filename in os.listdir(prov_path))

    def do_run(self):
        sim.setup(timestep=1.0, min_delay=1.0)
        sim.Population(10, sim.IF_curr_exp(), label='pop_1')
        sim.run(500)
        prov_path = SpynnakerDataView.get_app_provenance_dir_path()

        self.assertFalse(self.check_for_iobufs(prov_path))
        sim.end()
        self.assertFalse(self.check_for_iobufs(prov_path))

    def test_do_run(self):
        self.runsafe(self.do_run)

if __name__ == '__main__':
    x = TestNoIobufDuringRun()
    x.do_run()