"""Simplifies gold-miner output into true/false values to reduce numeric confusion"""

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter, FileType
from logging import debug, info, warning, error, critical
import logging
import re
import sys
import pyfsdb
import collections


def parse_args():
    parser = ArgumentParser(
        formatter_class=ArgumentDefaultsHelpFormatter,
        description=__doc__,
        epilog="Exmaple Usage: ",
    )

    parser.add_argument(
        "-g", "--gold", default=None, type=str, help="Protocol of interest"
    )

    parser.add_argument(
        "-t", "--time-bin-size", default=20, type=float, help="Time bin size"
    )

    parser.add_argument(
        "-f",
        "--output-fsdb",
        action="store_true",
        help="Output FSDB formatted complete dataset",
    )

    parser.add_argument(
        "--log-level",
        "--ll",
        default="info",
        help="Define the logging verbosity level (debug, info, warning, error, fotal, critical).",
    )

    parser.add_argument(
        "input_file",
        type=FileType("r"),
        nargs="?",
        default=sys.stdin,
        help="Input file to read generated by gold-miner",
    )

    parser.add_argument(
        "output_file",
        type=FileType("w"),
        nargs="?",
        default=sys.stdout,
        help="Where to write results",
    )

    args = parser.parse_args()
    log_level = args.log_level.upper()
    logging.basicConfig(level=log_level, format="%(levelname)-10s:\t%(message)s")
    return args


def output_results(
    timestamp,
    current_values,
    spi_counts,
    gold_target=None,
    output_handle=sys.stdout,
    output_fsdb=False,
):
    for spi in spi_counts:
        tokens = list(current_values[spi].keys())
        if len(tokens) <= 1:
            warning(f"Insufficient gold parameters identified for {timestamp=}")
            warning(f"Only keys found: {tokens}")
            warning(f"skipping")
            return
            # exit(1)  # may be overkill for low volumes?

        best_value = -100000
        best_token = None
        output_spi = re.sub("_+", "_", re.sub("[, '()]", "_", spi))

        for rank, token in enumerate(
            sorted(tokens, key=lambda x: current_values[spi][x])
        ):
            if current_values[spi][token] > best_value:
                best_value = current_values[spi][token]
                best_token = token
            if output_fsdb and spi_counts[spi] > 0:
                output_fsdb.append(
                    [
                        timestamp,
                        output_spi,
                        token,
                        current_values[spi][token] / spi_counts[spi],
                        token == gold_target,
                        rank,
                    ]
                )

        if gold_target:
            if best_token == gold_target:
                best_token = "True"
            else:
                best_token = "False"

        if output_fsdb:
            pass  # doing this above
        else:
            output_handle.write(
                f"{timestamp},{output_spi},{best_token}, {best_value}\n"
            )


def main():
    args = parse_args()

    inh = pyfsdb.Fsdb(
        file_handle=args.input_file,
        converters={"confidence": float, "timestamp": float},
        return_type=pyfsdb.RETURN_AS_DICTIONARY,
    )

    output_fsdb = False
    if args.output_fsdb:
        output_fsdb = pyfsdb.Fsdb(out_file_handle=args.output_file)
        output_fsdb.out_column_names = [
            "timestamp",
            "spi",
            "token",
            "average_value",
            "true",
            "rank",
        ]

    timestamp_counter = 1
    next_timestamp = None
    current_values = collections.defaultdict(dict)  # indexed by spi key
    spi_counts = collections.Counter()  # indexed by spi key

    for row in inh:
        if "identifier" in row:
            spi = row["identifier"]
        else:
            spi = row["spi"]
        token = row["token"]
        confidence = row["confidence"]
        timestamp = row["timestamp"]

        if next_timestamp is None:
            next_timestamp = timestamp + args.time_bin_size
        elif timestamp >= next_timestamp:
            debug(f"values:  {current_values}")
            debug(f"counts: {spi_counts}")

            next_timestamp += args.time_bin_size

            # report results
            output_results(
                timestamp_counter * args.time_bin_size,
                current_values,
                spi_counts,
                gold_target=args.gold,
                output_handle=args.output_file,
                output_fsdb=output_fsdb,
            )

            # reset these
            current_values = collections.defaultdict(dict)
            spi_counts = collections.Counter()
            timestamp_counter += 1

        if token not in current_values[spi]:
            current_values[spi][token] = 0.0
        current_values[spi][token] += confidence
        spi_counts[spi] += 1


if __name__ == "__main__":
    main()
