"""sonusai metrics_summary

usage: lsdb [-vlh] [-i MIXID] [-n NCPU] LOCATION

Options:
    -h, --help
    -v, --verbose
    -l, --write-list            Write .csv file list of all mixture metrics
    -i MIXID, --mixid MIXID     Mixture ID(s) to analyze. [default: *].
    -n, --num_process NCPU      Number of parallel processes to use [default: auto]

Summarize mixture metrics across a SonusAI mixture database where metrics have been generated by SonusAI genmetrics.

Inputs:
    LOCATION     A SonusAI mixture database directory with mixdb.db and pre-generated metrics from SonusAI genmetrics.

"""

import numpy as np
import pandas as pd

DB_99 = np.power(10, 99 / 10)
DB_N99 = np.power(10, -99 / 10)


def _process_mixture(
    m_id: int,
    location: str,
    all_metric_names: list[str],
    scalar_metric_names: list[str],
    string_metric_names: list[str],
    frame_metric_names: list[str],
    bin_metric_names: list[str],
    ptab_labels: list[str],
) -> tuple[pd.DataFrame, pd.DataFrame]:
    from os.path import basename

    from sonusai.constants import SAMPLE_RATE
    from sonusai.metrics import calc_wer
    from sonusai.mixture import MixtureDatabase

    mixdb = MixtureDatabase(location)

    # Process mixture
    # for mixid in mixids:
    samples = mixdb.mixture(m_id).samples
    duration = samples / SAMPLE_RATE
    tf_frames = mixdb.mixture_transform_frames(m_id)
    feat_frames = mixdb.mixture_feature_frames(m_id)
    mxsnr = mixdb.mixture(m_id).noise.snr
    ti = mixdb.mixture(m_id).sources["primary"].file_id
    ni = mixdb.mixture(m_id).noise.file_id
    t0file = basename(mixdb.source_file(ti).name)
    nfile = basename(mixdb.source_file(ni).name)

    all_metrics = mixdb.mixture_metrics(m_id, all_metric_names)

    # replace lists with first value (ignore mixup)
    scalar_metrics = {
        key: all_metrics[key][0] if isinstance(all_metrics[key], list) else all_metrics[key]
        for key in scalar_metric_names
    }
    string_metrics = {
        key: all_metrics[key][0] if isinstance(all_metrics[key], list) else all_metrics[key]
        for key in string_metric_names
    }

    # Convert strings into word count
    for key in string_metrics:
        string_metrics[key] = calc_wer(string_metrics[key], string_metrics[key]).words

    # Collect pandas table values note: must match given ptab_labels
    ptab_data: list = [
        mxsnr,
        *scalar_metrics.values(),
        *string_metrics.values(),
        tf_frames,
        duration,
        t0file,
        nfile,
    ]

    ptab1 = pd.DataFrame([ptab_data], columns=ptab_labels, index=[m_id])

    # TODO: collect frame metrics and bin metrics

    return ptab1, ptab1


def main() -> None:
    from docopt import docopt

    from . import __version__ as sai_version
    from .utils.docstring import trim_docstring

    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)

    verbose = args["--verbose"]
    wrlist = args["--write-list"]
    mixids = args["--mixid"]
    location = args["LOCATION"]
    num_proc = args["--num_process"]

    from functools import partial
    from os.path import basename
    from os.path import join

    import psutil

    from . import create_file_handler
    from . import initial_log_messages
    from . import logger
    from . import update_console_handler
    from .mixture.mixdb import MixtureDatabase
    from .utils.create_timestamp import create_timestamp
    from .utils.parallel import par_track
    from .utils.parallel import track

    mixdb = MixtureDatabase(location)
    print(f"Found SonusAI mixture database with {mixdb.num_mixtures} mixtures.")

    # Only check first and last mixture in order to save time
    metrics_present = mixdb.cached_metrics([0, mixdb.num_mixtures - 1])  # return pre-generated metrics in mixdb tree
    if "mxsnr" in metrics_present:
        metrics_present.remove("mxsnr")

    num_metrics_present = len(metrics_present)
    if num_metrics_present < 1:
        print(f"mixdb reports no pre-generated metrics are present. Nothing to summarize in {location}, exiting ...")
        return

    # Setup logging file
    timestamp = create_timestamp()  # string good for embedding into filenames
    mixdb_fname = basename(location)
    if verbose:
        create_file_handler(join(location, "metrics_summary.log"))
        update_console_handler(verbose)
        initial_log_messages("metrics_summary")
        logger.info(f"Logging summary of SonusAI mixture database at {location}")
    else:
        update_console_handler(verbose)

    logger.info("")
    mixids = mixdb.mixids_to_list(mixids)
    if len(mixids) < mixdb.num_mixtures:
        logger.info(
            f"Processing a subset of {len(mixids)} out of total mixdb mixtures of {mixdb.num_mixtures}, "
            f"summary results will not include entire dataset."
        )
        fsuffix = f"_s{len(mixids)}t{mixdb.num_mixtures}"
    else:
        logger.info(
            f"Summarizing SonusAI mixture database with {mixdb.num_mixtures} mixtures "
            f"and {num_metrics_present} pre-generated metrics ..."
        )
        fsuffix = ""

    metric_sup = mixdb.supported_metrics
    ft_bins = mixdb.ft_config.bin_end - mixdb.ft_config.bin_start + 1  # bins of forward transform
    # Pre-process first mixid to gather metrics into 4 types: scalar, str (scalar word cnt), frame-array, bin-array
    # Collect list of indices for each
    scalar_metric_names: list[str] = []
    string_metric_names: list[str] = []
    frame_metric_names: list[str] = []
    bin_metric_names: list[str] = []
    all_metrics = mixdb.mixture_metrics(mixids[0], metrics_present)
    tf_frames = mixdb.mixture_transform_frames(mixids[0])
    for metric in metrics_present:
        metval = all_metrics[metric]  # get metric value
        logger.debug(f"First mixid {mixids[0]} metric {metric} = {metval}")
        if isinstance(metval, list):
            if len(metval) > 1:
                logger.warning(f"Mixid {mixids[0]} metric {metric} has a list with more than 1 element, using first.")
            metval = metval[0]  # remove any list
        if isinstance(metval, float | int):
            logger.debug(f"Metric is scalar {type(metval)}, entering in summary table.")
            scalar_metric_names.append(metric)
        elif isinstance(metval, str):
            logger.debug("Metric is string, will summarize with word count.")
            string_metric_names.append(metric)
        elif isinstance(metval, np.ndarray):
            if metval.ndim == 1:
                if metval.size == tf_frames:
                    logger.debug("Metric is frames vector.")
                    frame_metric_names.append(metric)
                elif metval.size == ft_bins:
                    logger.debug("Metric is bins vector.")
                    bin_metric_names.append(metric)
                else:
                    logger.warning(f"Mixid {mixids[0]} metric {metric} is a vector of improper size, ignoring.")

    # Setup pandas table for summarizing scalar metrics, always include mxsnr first
    ptab_labels = [
        "mxsnr",
        *scalar_metric_names,
        *string_metric_names,
        "fcnt",
        "duration",
        "t0file",
        "nfile",
    ]

    num_cpu = psutil.cpu_count()
    cpu_percent = psutil.cpu_percent(interval=1)
    logger.info("")
    logger.info(f"#CPUs: {num_cpu}, current CPU utilization: {cpu_percent}%")
    logger.info(f"Memory utilization: {psutil.virtual_memory().percent}%")
    if num_proc == "auto":
        use_cpu = int(num_cpu * (0.9 - cpu_percent / 100))  # default use 80% of available cpus
    elif num_proc == "None":
        use_cpu = None
    else:
        use_cpu = min(max(int(num_proc), 1), num_cpu)

    logger.info(f"Summarizing metrics for {len(mixids)} mixtures using {use_cpu} parallel processes")

    # progress = tqdm(total=len(mixids), desc='calc_metric_spenh', mininterval=1)
    progress = track(total=len(mixids))
    if use_cpu is None:
        no_par = True
        num_cpus = None
    else:
        no_par = False
        num_cpus = use_cpu

    all_metrics_tables = par_track(
        partial(
            _process_mixture,
            location=location,
            all_metric_names=metrics_present,
            scalar_metric_names=scalar_metric_names,
            string_metric_names=string_metric_names,
            frame_metric_names=frame_metric_names,
            bin_metric_names=bin_metric_names,
            ptab_labels=ptab_labels,
        ),
        mixids,
        progress=progress,
        num_cpus=num_cpus,
        no_par=no_par,
    )
    progress.close()

    # Done with mixtures, write out summary metrics
    header_args = {
        "mode": "a",
        "encoding": "utf-8",
        "index": False,
        "header": False,
    }
    table_args = {
        "mode": "a",
        "encoding": "utf-8",
    }
    ptab1 = pd.concat([item[0] for item in all_metrics_tables])
    if wrlist:
        wlcsv_name = str(join(location, "metric_summary_list" + fsuffix + ".csv"))
        pd.DataFrame([["Timestamp", timestamp]]).to_csv(wlcsv_name, header=False, index=False)
        pd.DataFrame([f"Metric list for {mixdb_fname}:"]).to_csv(wlcsv_name, mode="a", header=False, index=False)
        ptab1.round(2).to_csv(wlcsv_name, **table_args)
    ptab1_sorted = ptab1.sort_values(by=["mxsnr", "t0file"])

    # Create metrics table except -99 SNR
    ptab1_nom99 = ptab1_sorted[ptab1_sorted.mxsnr != -99]

    # Create summary by SNR for all scalar metrics, taking mean
    mtab_snr_summary = None
    for snri in range(0, len(mixdb.snrs)):
        tmp = ptab1_sorted.query("mxsnr==" + str(mixdb.snrs[snri])).mean(numeric_only=True).to_frame().T
        # avoid nan when subset of mixids specified (i.e. no mixtures exist for an SNR)
        if ~np.isnan(tmp.iloc[0].to_numpy()[0]).any():
            mtab_snr_summary = pd.concat([mtab_snr_summary, tmp])
    mtab_snr_summary = mtab_snr_summary.sort_values(by=["mxsnr"], ascending=False)

    # Write summary to .csv
    snrcsv_name = str(join(location, "metric_summary_snr" + fsuffix + ".csv"))
    nmix = len(mixids)
    nmixtot = mixdb.num_mixtures
    pd.DataFrame([["Timestamp", timestamp]]).to_csv(snrcsv_name, header=False, index=False)
    pd.DataFrame(['"Metrics avg over each SNR:"']).to_csv(snrcsv_name, **header_args)
    mtab_snr_summary.round(2).T.to_csv(snrcsv_name, index=True, header=False, mode="a", encoding="utf-8")
    pd.DataFrame(["--"]).to_csv(snrcsv_name, header=False, index=False, mode="a")
    pd.DataFrame([f'"Metrics stats over {nmix} mixtures out of {nmixtot} total:"']).to_csv(snrcsv_name, **header_args)
    ptab1.describe().round(2).T.to_csv(snrcsv_name, index=True, **table_args)
    pd.DataFrame(["--"]).to_csv(snrcsv_name, header=False, index=False, mode="a")
    pd.DataFrame([f'"Metrics stats over {len(ptab1_nom99)} non -99db mixtures out of {nmixtot} total:"']).to_csv(
        snrcsv_name, **header_args
    )
    ptab1_nom99.describe().round(2).T.to_csv(snrcsv_name, index=True, **table_args)

    # Write summary to text file
    snrtxt_name = str(join(location, "metric_summary_snr" + fsuffix + ".txt"))
    with open(snrtxt_name, "w") as f:
        print(f"Timestamp: {timestamp}", file=f)
        print("Metrics avg over each SNR:", file=f)
        print(
            mtab_snr_summary.round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True, header=False), file=f
        )
        print("", file=f)
        print(f"Metrics stats over {len(mixids)} mixtures out of {mixdb.num_mixtures} total:", file=f)
        print(ptab1.describe().round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True), file=f)
        print("", file=f)
        print(f"Metrics stats over {len(ptab1_nom99)} non -99db mixtures out of {mixdb.num_mixtures} total:", file=f)
        print(ptab1_nom99.describe().round(2).T.to_string(float_format=lambda x: f"{x:.2f}", index=True), file=f)


if __name__ == "__main__":
    from sonusai import exception_handler
    from sonusai.utils import register_keyboard_interrupt

    register_keyboard_interrupt()
    try:
        main()
    except Exception as e:
        exception_handler(e)
