"""
Error analysis of the false positive and false negative breasts.

    python scripts/validation/analyze_false_classification.py <oof_dir> <metadata_file>

Reads the per case out-of-fold predictions of predict_oof.py, turns the scores into decisions, and
tabulates the errors by BI-RADS category and by prior biopsy status, per model.

The threshold is the one used throughout the paper: within each outer fold, the smallest threshold
that still reaches --operating_point specificity. It has to be found inside a fold because the five
folds come from separately trained models whose score scales are not comparable, and it is taken
from the same fold it is applied to, so the absolute error counts are optimistic. What the tables
are for is the *composition* of the errors, which is far less sensitive to that.

Per stratum the table reports, for each model:

    n           breasts in the stratum
    malignant   how many of them are malignant
    FP, FPR     false positives among the benign breasts of the stratum, and the rate
    FN, FNR     false negatives among the malignant breasts of the stratum, and the rate

so a row says how the errors are distributed over the stratum, not how often the stratum occurs.

BI-RADS and prior biopsy are read per side from the metadata export via the configured column
names (`birads_left`/`birads_right`, `biopsy_left`/`biopsy_right`). The biopsy columns are empty in
the default config because metadata_cleaned.csv does not carry that field; point them at real
columns in the site config and the biopsy tables appear. Any further per-side column can be added
with --extra_stratum <config key>.

Writes a report, one CSV per stratification, and the individual misclassified breasts as
false_cases.csv so they can be reviewed.
"""

import argparse
from pathlib import Path

import numpy as np
import pandas as pd

from auto_detect_breast_mri.config import normalise_indication_code, resolve_path
from auto_detect_breast_mri.evaluation.cluster_bootstrap import sensitivity_at_specificity
from auto_detect_breast_mri.evaluation.oof_table import (available_fractions, load_patient_map,
                                                         load_side_map)

DEFAULT_OPERATING_TARGET = 0.9
MISSING_STRATUM = 'unknown'
# stratifications tried by default: config key -> column title in the report
DEFAULT_STRATA = {'birads': 'BI-RADS', 'biopsy': 'prior biopsy'}


def read_predictions(oof_dir, metadata_file, fraction=None):
    """
    Long table of per breast predictions, one row per (model, protocol, fold, breast).

    :return: (DataFrame, dict describing what was read)
    """
    from auto_detect_breast_mri.evaluation.oof_table import FILE_PATTERN

    per_fraction = available_fractions(oof_dir)
    if not per_fraction:
        raise FileNotFoundError(f"No oof_*_fold*.csv in {oof_dir}. Run predict_oof.py first.")
    chosen = max(per_fraction) if fraction is None else float(fraction)
    if chosen not in per_fraction:
        raise ValueError(f"{oof_dir} has no predictions for fraction {chosen:g}; it holds "
                         f"{', '.join(f'{value:g}' for value in sorted(per_fraction))}.")
    if fraction is None and len(per_fraction) > 1:
        print(f"  {oof_dir} holds {len(per_fraction)} training fractions; using {chosen:g}.")

    frames = []
    for path in sorted(per_fraction[chosen]):
        parsed = FILE_PATTERN.match(Path(path).name)
        frame = pd.read_csv(path, dtype={'examination_id': str, 'side': str})
        frame['architecture'] = parsed.group('architecture')
        frames.append(frame)
    table = pd.concat(frames, ignore_index=True)
    table['side'] = table['side'].str.lower()
    table['model'] = table['architecture'] + '_' + table['protocol']

    patient_map = load_patient_map(metadata_file, table['examination_id'].unique())
    table['patient_id'] = table['examination_id'].map(patient_map)
    report = {'fraction': chosen, 'models': sorted(table['model'].unique()),
              'without_patient_id': int(table['patient_id'].isna().sum())}
    report['breasts'] = int(table.groupby('model').size().max())
    return table, report


def classify(table, target=DEFAULT_OPERATING_TARGET):
    """
    Turn scores into decisions, with the threshold found per (model, fold).

    :return: (table with 'predicted' and 'outcome' columns, dict (model, fold) -> threshold)
    """
    table = table.copy()
    table['predicted'] = np.nan
    thresholds = {}
    for (model, fold), part in table.groupby(['model', 'outer_fold']):
        _, _, threshold = sensitivity_at_specificity(part['label'].to_numpy(),
                                                     part['score'].to_numpy(), target)
        thresholds[(model, int(fold))] = float(threshold)
        table.loc[part.index, 'predicted'] = (part['score'] >= threshold).astype(int)
    table['predicted'] = table['predicted'].astype(int)
    table['outcome'] = np.select(
        [(table['label'] == 1) & (table['predicted'] == 1),
         (table['label'] == 0) & (table['predicted'] == 1),
         (table['label'] == 1) & (table['predicted'] == 0)],
        ['TP', 'FP', 'FN'], default='TN')
    return table, thresholds


def attach_strata(table, metadata_file, keys):
    """
    Add one column per stratification, read per side from the metadata.

    :param keys: config key prefixes, e.g. ['birads', 'biopsy']
    :return: (table, dict key -> note about what happened)
    """
    table = table.copy()
    notes = {}
    identifiers = table['examination_id'].unique()
    for key in keys:
        # '2.0' and 2 are one category; normalise_indication_code strips the trailing .0
        mapping, report = load_side_map(metadata_file, key, identifiers,
                                        normalise=normalise_indication_code)
        if mapping is None:
            notes[key] = (f"not available: {', '.join(report['missing_columns'])}. Configure "
                          f"columns.{key}_left / columns.{key}_right to enable this table.")
            continue
        pairs = list(zip(table['examination_id'], table['side']))
        table[key] = [mapping.get(pair, MISSING_STRATUM) for pair in pairs]
        unknown = int((table[key] == MISSING_STRATUM).sum() / max(1, table['model'].nunique()))
        notes[key] = (f"read from the metadata; {unknown} breasts per model have no value "
                      f"and are reported as '{MISSING_STRATUM}'.")
    return table, notes


def _stratum_order(value):
    """Numeric strata sort numerically, everything else (e.g. 'unknown') after them."""
    try:
        return (0, float(value), '')
    except (TypeError, ValueError):
        return (1, 0.0, str(value))


def tabulate(table, key):
    """
    False positives and negatives per stratum and model.

    :return: DataFrame, one row per (model, stratum)
    """
    rows = []
    for model in sorted(table['model'].unique()):
        part = table[table['model'] == model]
        for stratum in sorted(part[key].unique(), key=_stratum_order):
            cases = part[part[key] == stratum]
            benign, malignant = cases[cases['label'] == 0], cases[cases['label'] == 1]
            false_positive = int((benign['outcome'] == 'FP').sum())
            false_negative = int((malignant['outcome'] == 'FN').sum())
            rows.append({
                'model': model, key: stratum, 'n': len(cases),
                'benign': len(benign), 'malignant': len(malignant),
                'FP [n]': false_positive,
                'FPR': false_positive / len(benign) if len(benign) else np.nan,
                'FN [n]': false_negative,
                'FNR': false_negative / len(malignant) if len(malignant) else np.nan,
                'FP [%]': np.nan, 'FN [%]': np.nan})
        total_fp = sum(row['FP [n]'] for row in rows if row['model'] == model)
        total_fn = sum(row['FN [n]'] for row in rows if row['model'] == model)
        for row in rows:
            if row['model'] == model:
                row['FP [%]'] = row['FP [n]'] / total_fp if total_fp else np.nan
                row['FN [%]'] = row['FN [n]'] / total_fn if total_fn else np.nan
    return pd.DataFrame(rows)


def format_table(frame, key, title, decimals=3):
    lines = [f"{title}", "-" * 100,
             f"{'model':<16} {str(key):>12} {'n':>6} {'benign':>7} {'malig':>6} "
             f"{'FP':>5} {'FPR':>7} {'%ofFP':>7} {'FN':>5} {'FNR':>7} {'%ofFN':>7}"]
    for model, block in frame.groupby('model', sort=True):
        for _, row in block.iterrows():
            def number(value, width=7):
                return f"{value:>{width}.{decimals}f}" if not pd.isna(value) else f"{'-':>{width}}"
            lines.append(f"{row['model']:<16} {str(row[key]):>12} {row['n']:>6} "
                         f"{row['benign']:>7} {row['malignant']:>6} {row['FP [n]']:>5} "
                         f"{number(row['FPR'])} {number(row['FP [%]'])} "
                         f"{row['FN [n]']:>5} {number(row['FNR'])} {number(row['FN [%]'])}")
        lines.append("")
    return "\n".join(lines)


def format_report(overall, tables, notes, strata_notes, report, target, thresholds, decimals=3):
    lines = []
    add = lines.append
    add("=" * 100)
    add("Error analysis of the false positive and false negative breasts")
    add("=" * 100)
    add("")
    add(f"{'models':<26}{', '.join(report['models'])}")
    add(f"{'training fraction':<26}{report['fraction']:g}")
    add(f"{'breasts per model':<26}{report['breasts']}")
    add(f"{'threshold':<26}per fold, smallest reaching {target:.0%} specificity")
    for key, note in strata_notes.items():
        if 'no value' in note:
            add(f"{'without ' + key:<26}{note.split(';')[1].strip().split(' breasts')[0]} breasts "
                f"per model, reported as '{MISSING_STRATUM}'")
    add("")
    add("-" * 100)
    add("0. Confusion counts per model, at that threshold")
    add("-" * 100)
    add(f"{'model':<16} {'TP':>6} {'FP':>6} {'FN':>6} {'TN':>6} {'sensitivity':>12} "
        f"{'specificity':>12}")
    for _, row in overall.iterrows():
        add(f"{row['model']:<16} {row['TP']:>6} {row['FP']:>6} {row['FN']:>6} {row['TN']:>6} "
            f"{row['sensitivity']:>12.{decimals}f} {row['specificity']:>12.{decimals}f}")
    add("")
    for key, title in tables:
        add("-" * 100)
        add(title)
        add("-" * 100)
        add(notes[key])
        add("")
    add("-" * 100)
    add("How to read this")
    add("-" * 100)
    add("FPR is the false positive rate among the *benign* breasts of that stratum, FNR the false")
    add("negative rate among its *malignant* breasts, so the two rates never share a denominator.")
    add("'%ofFP' and '%ofFN' say how the model's errors are distributed over the strata and sum")
    add("to 1 per model. A stratum with few cases can have an extreme rate on one or two breasts.")
    add("")
    add("The threshold is taken from the same fold it is applied to, so the absolute counts are")
    add("optimistic. The composition of the errors across strata is the point here, and it is far")
    add("less sensitive to the threshold than the totals are.")
    for key, note in strata_notes.items():
        add("")
        add(f"{key}: {note}")
    return "\n".join(lines)


def build_parser():
    parser = argparse.ArgumentParser(prog="analyze_false_classification", description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("oof_dir", help="Folder holding oof_<arch>_fold<k>[_frac=<f>].csv")
    parser.add_argument("metadata_file", help="Metadata export with the per-side columns.")
    parser.add_argument("-o", "--output_path", default=None,
                        help="Folder the report and the CSVs are written to.")
    parser.add_argument("--fraction", type=float, default=None,
                        help="Training fraction to read. Default: the largest present.")
    parser.add_argument("--operating_point", type=float, default=DEFAULT_OPERATING_TARGET,
                        help="Specificity the threshold is set at. Default: %(default)s")
    parser.add_argument("--extra_stratum", nargs='+', default=[], metavar="KEY",
                        help="Further per-side config column keys to tabulate by, e.g. indication.")
    parser.add_argument("-d", "--decimals", type=int, default=3)
    return parser


def main():
    args = build_parser().parse_args()

    output_dir = Path(resolve_path(args.output_path, "output_root", "output folder")) / "errors"
    output_dir.mkdir(parents=True, exist_ok=True)

    table, report = read_predictions(args.oof_dir, args.metadata_file, args.fraction)
    table, thresholds = classify(table, args.operating_point)

    keys = list(DEFAULT_STRATA) + [key for key in args.extra_stratum if key not in DEFAULT_STRATA]
    table, strata_notes = attach_strata(table, args.metadata_file, keys)

    counts = (table.groupby(['model', 'outcome']).size().unstack(fill_value=0)
              .reindex(columns=['TP', 'FP', 'FN', 'TN'], fill_value=0).reset_index())
    counts['sensitivity'] = counts['TP'] / (counts['TP'] + counts['FN']).replace(0, np.nan)
    counts['specificity'] = counts['TN'] / (counts['TN'] + counts['FP']).replace(0, np.nan)
    counts.to_csv(output_dir / "confusion_counts.csv", index=False)

    tables, notes = [], {}
    for key in keys:
        title = f"Errors by {DEFAULT_STRATA.get(key, key)}"
        if key not in table.columns:
            notes[key] = f"NOT AVAILABLE. {strata_notes[key]}"
            tables.append((key, title))
            continue
        frame = tabulate(table, key)
        frame.to_csv(output_dir / f"false_classification_by_{key}.csv", index=False)
        notes[key] = format_table(frame, key, "", args.decimals)
        tables.append((key, title))

    false_cases = table[table['outcome'].isin(['FP', 'FN'])].copy()
    columns = ['model', 'outer_fold', 'examination_id', 'side', 'label', 'score', 'outcome']
    columns += [key for key in keys if key in table.columns]
    false_cases[columns].sort_values(['model', 'outcome', 'examination_id']).to_csv(
        output_dir / "false_cases.csv", index=False)

    text = format_report(counts, tables, notes, strata_notes, report, args.operating_point,
                         thresholds, args.decimals)
    (output_dir / "false_classification_report.txt").write_text(text)
    print()
    print(text)
    print(f"\n{len(false_cases)} misclassified breast-model rows written to "
          f"{output_dir / 'false_cases.csv'}")
    print(f"Wrote the report and the CSVs to {output_dir}")


if __name__ == '__main__':
    main()