assignment_misdirection_analyzer.py

#!/usr/bin/env python3
"""
Unified Assignment Misdirection Analyzer

This script supports two input modes using the SAME CLI shape:

Mode A: Post-evaluation JSON mode
---------------------------------
Input files already contain:
  evaluation.relation_metrics.fp_relations
  evaluation.relation_metrics.fn_relations
  evaluation.relation_metrics.tp_relations (optional)

This is the current assignment_misdirection_analyzer.py behavior.
Typical use: *_wo_legend path_generation JSONs.

Mode B: Raw prediction JSON mode
--------------------------------
Input files contain predicted graph edges only:
  nodes: [...]
  edges: [...]

The script automatically pairs each prediction file to a GT file and computes:
  - assign_gt / assign_pred / assign_fn / assign_fp
  - assign_misdirected
  - assign_missing_not_misdirected

Typical use: with-legend SAG_with_legend / MAG_with_legend / DAG_with_legend directories.

Outputs
-------
- <out_prefix>_per_sample.csv
- <out_prefix>_summary.json
- <out_prefix>_table.md

CLI
---
python3 assignment_misdirection_analyzer.py \
  --dir /path/to/input_dir \
  --glob "*path_generation.json" \
  --out_prefix assign_misdirection

Optional:
  --out_dir /path/to/write/results
  --gt_root /Users/.../BlueSkyEvaluation/gt

If --out_dir is omitted, outputs are written into --dir.
If --gt_root is omitted, a sensible default path is used.
"""

from __future__ import annotations

import argparse
import csv
import glob
import json
import os
import re
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Tuple, Set
from collections import Counter, defaultdict


# -----------------------
# Normalization utilities
# -----------------------

def norm(s: str) -> str:
    if s is None:
        return ""
    s = str(s).strip().lower().replace("_", " ")
    return " ".join(s.split())


def safe_load_json(path: str) -> Optional[dict]:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


# -----------------------
# Type / category config
# -----------------------

TYPE_USER = "user_attributes"
TYPE_OBJ = "object_attributes"
TYPE_PC = "policy_classes"

TRACKED_PAIRS = {
    (TYPE_USER, TYPE_USER),
    (TYPE_USER, TYPE_PC),
    (TYPE_OBJ, TYPE_OBJ),
    (TYPE_OBJ, TYPE_PC),
}

REL_ASSIGN = "assign"
REL_ASSOC = "associate"
REL_PROHIBIT = "prohibit"
REL_PERMIT = "permit"

REL_ALIASES = {
    "assign": REL_ASSIGN,
    "assignment": REL_ASSIGN,
    "associate": REL_ASSOC,
    "association": REL_ASSOC,
    "prohibit": REL_PROHIBIT,
    "prohibition": REL_PROHIBIT,
    "deny": REL_PROHIBIT,
    "permit": REL_PERMIT,
}


def canon_rel(r: str) -> str:
    return REL_ALIASES.get(norm(r), norm(r))


def pair_key(src_t: str, dst_t: str) -> str:
    return f"{src_t} -> {dst_t}"


# -----------------------
# Data structure
# -----------------------

@dataclass
class SampleRow:
    sample_file: str
    assign_gt: int
    assign_pred: int
    assign_fn: int
    assign_fp: int
    assign_misdirected: int
    assign_missing_not_misdirected: int
    unknown_endpoints: int

    uu_gt: int; uu_pred: int; uu_fn: int; uu_fp: int; uu_mis: int
    up_gt: int; up_pred: int; up_fn: int; up_fp: int; up_mis: int
    oo_gt: int; oo_pred: int; oo_fn: int; oo_fp: int; oo_mis: int
    op_gt: int; op_pred: int; op_fn: int; op_fp: int; op_mis: int


# -----------------------
# Shared helpers
# -----------------------

def build_type_lookup(nodes: List[dict]) -> Dict[str, str]:
    lut = {}
    for n in nodes or []:
        if not isinstance(n, dict):
            continue
        t = n.get("type")
        c = n.get("content")
        if t and c:
            lut[norm(c)] = t
    return lut


def classify_pair(
    type_lut: Dict[str, str], src_name: str, dst_name: str
) -> Tuple[Optional[Tuple[str, str]], int]:
    s = type_lut.get(norm(src_name))
    d = type_lut.get(norm(dst_name))
    if s is None or d is None:
        return None, 1
    return (s, d), 0


def safe_div(a: int, b: int) -> Optional[float]:
    return (a / b) if b > 0 else None


def fmt(x: Optional[float], nd: int = 3) -> str:
    return "n/a" if x is None else f"{x:.{nd}f}"


# -----------------------
# Mode detection
# -----------------------

def is_post_eval_json(j: dict) -> bool:
    relm = (j.get("evaluation", {}) or {}).get("relation_metrics", {}) or {}
    return any(k in relm for k in ("fp_relations", "fn_relations", "tp_relations"))


def is_raw_pred_json(j: dict) -> bool:
    return isinstance(j.get("nodes", None), list) and isinstance(j.get("edges", None), list)


# -----------------------
# Mode A: post-eval JSON analysis
# -----------------------

def get_rel_lists(j: dict) -> Tuple[List[list], List[list], List[list]]:
    relm = (j.get("evaluation", {}) or {}).get("relation_metrics", {}) or {}
    tp = relm.get("tp_relations", []) or []
    fp = relm.get("fp_relations", []) or []
    fn = relm.get("fn_relations", []) or []
    tp = [r for r in tp if isinstance(r, list) and len(r) >= 3]
    fp = [r for r in fp if isinstance(r, list) and len(r) >= 3]
    fn = [r for r in fn if isinstance(r, list) and len(r) >= 3]
    return tp, fp, fn


def is_assign(rel: list) -> bool:
    return norm(rel[2]) == REL_ASSIGN


def analyze_post_eval_file(path: str) -> Tuple[SampleRow, Counter]:
    with open(path, "r", encoding="utf-8") as f:
        j = json.load(f)

    type_lut = build_type_lookup(j.get("nodes", []))
    tp_rel, fp_rel, fn_rel = get_rel_lists(j)

    tp_assign = [r for r in tp_rel if is_assign(r)]
    fp_assign = [r for r in fp_rel if is_assign(r)]
    fn_assign = [r for r in fn_rel if is_assign(r)]

    fp_assign_set = {(norm(r[0]), norm(r[1])) for r in fp_assign}

    counts = defaultdict(lambda: {"gt": 0, "pred": 0, "fn": 0, "fp": 0, "mis": 0})
    unknown_endpoints = 0

    def process_list(rel_list: List[list], kind: str):
        nonlocal unknown_endpoints
        for r in rel_list:
            src, dst = r[0], r[1]
            pair, unk = classify_pair(type_lut, src, dst)
            unknown_endpoints += unk
            if pair is None or pair not in TRACKED_PAIRS:
                continue
            key = pair_key(*pair)
            if kind == "tp":
                counts[key]["gt"] += 1
                counts[key]["pred"] += 1
            elif kind == "fn":
                counts[key]["gt"] += 1
                counts[key]["fn"] += 1
            elif kind == "fp":
                counts[key]["pred"] += 1
                counts[key]["fp"] += 1

    process_list(tp_assign, "tp")
    process_list(fn_assign, "fn")
    process_list(fp_assign, "fp")

    misdirected_total = 0
    missing_not_misdirected_total = 0
    misdirected_edges_counter = Counter()

    for r in fn_assign:
        src_n, dst_n = norm(r[0]), norm(r[1])
        pair, _ = classify_pair(type_lut, r[0], r[1])
        if pair is None or pair not in TRACKED_PAIRS:
            continue

        key = pair_key(*pair)
        reversed_exists = (dst_n, src_n) in fp_assign_set
        if reversed_exists:
            counts[key]["mis"] += 1
            misdirected_total += 1
            misdirected_edges_counter[(src_n, dst_n)] += 1
        else:
            missing_not_misdirected_total += 1

    uu = counts[pair_key(TYPE_USER, TYPE_USER)]
    up = counts[pair_key(TYPE_USER, TYPE_PC)]
    oo = counts[pair_key(TYPE_OBJ, TYPE_OBJ)]
    op = counts[pair_key(TYPE_OBJ, TYPE_PC)]

    row = SampleRow(
        sample_file=os.path.basename(path),
        assign_gt=sum(v["gt"] for v in counts.values()),
        assign_pred=sum(v["pred"] for v in counts.values()),
        assign_fn=sum(v["fn"] for v in counts.values()),
        assign_fp=sum(v["fp"] for v in counts.values()),
        assign_misdirected=misdirected_total,
        assign_missing_not_misdirected=missing_not_misdirected_total,
        unknown_endpoints=unknown_endpoints,

        uu_gt=uu["gt"], uu_pred=uu["pred"], uu_fn=uu["fn"], uu_fp=uu["fp"], uu_mis=uu["mis"],
        up_gt=up["gt"], up_pred=up["pred"], up_fn=up["fn"], up_fp=up["fp"], up_mis=up["mis"],
        oo_gt=oo["gt"], oo_pred=oo["pred"], oo_fn=oo["fn"], oo_fp=oo["fp"], oo_mis=oo["mis"],
        op_gt=op["gt"], op_pred=op["pred"], op_fn=op["fn"], op_fp=op["fp"], op_mis=op["mis"],
    )
    return row, misdirected_edges_counter


# -----------------------
# Mode B: raw pred + GT analysis
# -----------------------

def canon_key(path: str, is_pred: bool) -> str:
    base = os.path.splitext(os.path.basename(path))[0]
    if is_pred:
        base = re.sub(r"_path_generation$", "", base)
        base = re.sub(r"_labeled_[a-z0-9]+$", "", base)
        base = re.sub(r"_labeled$", "", base)
    return base


def infer_subset_from_dir(input_dir: str) -> str:
    base = os.path.basename(os.path.normpath(input_dir))
    if base.endswith("_wo_legend"):
        base = base[:-10]
    return base


def infer_gt_dir(input_dir: str, gt_root: str) -> str:
    subset = infer_subset_from_dir(input_dir)
    return os.path.join(gt_root, subset)


def gt_extract_edges_and_types(gt: dict) -> Tuple[List[Tuple[str, str, str]], Dict[str, str]]:
    edges: List[Tuple[str, str, str]] = []
    type_lut: Dict[str, str] = {}

    pe = gt.get("policy_elements", {}) or {}

    for label in pe.get("user_attributes", []) or []:
        type_lut[norm(label)] = TYPE_USER
    for label in pe.get("object_attributes", []) or []:
        type_lut[norm(label)] = TYPE_OBJ

    pc = pe.get("policy_classes", None)
    if isinstance(pc, str) and pc:
        type_lut[norm(pc)] = TYPE_PC
    elif isinstance(pc, list):
        for label in pc:
            type_lut[norm(label)] = TYPE_PC

    for _, v in (gt.get("assignments", {}) or {}).items():
        if isinstance(v, dict) and "from" in v and "to" in v:
            edges.append((norm(v["from"]), norm(v["to"]), REL_ASSIGN))

    for _, v in (gt.get("associations", {}) or {}).items():
        if isinstance(v, dict) and "from" in v and "to" in v:
            edges.append((norm(v["from"]), norm(v["to"]), REL_ASSOC))

    for _, v in (gt.get("prohibitions", {}) or {}).items():
        if isinstance(v, dict) and "from" in v and "to" in v:
            edges.append((norm(v["from"]), norm(v["to"]), REL_PROHIBIT))

    for key in ("permits", "permissions", "permissions_relations", "permissions_edges"):
        block = gt.get(key, None)
        if isinstance(block, dict):
            for _, v in block.items():
                if isinstance(v, dict) and "from" in v and "to" in v:
                    edges.append((norm(v["from"]), norm(v["to"]), REL_PERMIT))
        elif isinstance(block, list):
            for v in block:
                if isinstance(v, dict) and "from" in v and "to" in v:
                    edges.append((norm(v["from"]), norm(v["to"]), REL_PERMIT))

    return edges, type_lut


def pred_extract_edges_and_types(pred: dict) -> Tuple[List[Tuple[str, str, str]], Dict[str, str]]:
    type_lut = build_type_lookup(pred.get("nodes", []))
    edges: List[Tuple[str, str, str]] = []
    for e in pred.get("edges", []) or []:
        if not isinstance(e, dict):
            continue
        src = e.get("source_name", e.get("from"))
        dst = e.get("target_name", e.get("to"))
        rel = e.get("relationship")
        if src and dst and rel:
            edges.append((norm(src), norm(dst), canon_rel(rel)))
    return edges, type_lut


def filter_assign(edges: List[Tuple[str, str, str]]) -> List[Tuple[str, str, str]]:
    return [e for e in edges if e[2] == REL_ASSIGN]


def analyze_raw_pair(gt_path: str, pred_path: str) -> Tuple[SampleRow, Counter]:
    gt = safe_load_json(gt_path)
    pr = safe_load_json(pred_path)
    if gt is None or pr is None:
        raise ValueError(f"Could not load pair: gt={gt_path}, pred={pred_path}")

    gt_edges, gt_types = gt_extract_edges_and_types(gt)
    pr_edges, pr_types = pred_extract_edges_and_types(pr)

    gt_assign = filter_assign(gt_edges)
    pr_assign = filter_assign(pr_edges)

    gt_assign_set = {(a, b) for (a, b, _) in gt_assign}
    pr_assign_set = {(a, b) for (a, b, _) in pr_assign}

    tp_pairs = gt_assign_set & pr_assign_set
    fn_pairs = gt_assign_set - pr_assign_set
    fp_pairs = pr_assign_set - gt_assign_set

    counts = defaultdict(lambda: {"gt": 0, "pred": 0, "fn": 0, "fp": 0, "mis": 0})
    unknown_endpoints = 0
    misdirected_total = 0
    missing_not_misdirected_total = 0
    misdirected_edges_counter = Counter()

    def add_pair(src: str, dst: str, kind: str):
        nonlocal unknown_endpoints
        pair, unk = classify_pair(gt_types, src, dst)
        unknown_endpoints += unk
        if pair is None or pair not in TRACKED_PAIRS:
            return
        key = pair_key(*pair)
        if kind == "tp":
            counts[key]["gt"] += 1
            counts[key]["pred"] += 1
        elif kind == "fn":
            counts[key]["gt"] += 1
            counts[key]["fn"] += 1
        elif kind == "fp":
            counts[key]["pred"] += 1
            counts[key]["fp"] += 1

    for a, b in tp_pairs:
        add_pair(a, b, "tp")
    for a, b in fn_pairs:
        add_pair(a, b, "fn")
    for a, b in fp_pairs:
        add_pair(a, b, "fp")

    for a, b in fn_pairs:
        pair, _ = classify_pair(gt_types, a, b)
        if pair is None or pair not in TRACKED_PAIRS:
            continue
        key = pair_key(*pair)
        if (b, a) in fp_pairs:
            counts[key]["mis"] += 1
            misdirected_total += 1
            misdirected_edges_counter[(a, b)] += 1
        else:
            missing_not_misdirected_total += 1

    uu = counts[pair_key(TYPE_USER, TYPE_USER)]
    up = counts[pair_key(TYPE_USER, TYPE_PC)]
    oo = counts[pair_key(TYPE_OBJ, TYPE_OBJ)]
    op = counts[pair_key(TYPE_OBJ, TYPE_PC)]

    row = SampleRow(
        sample_file=os.path.basename(pred_path),
        assign_gt=sum(v["gt"] for v in counts.values()),
        assign_pred=sum(v["pred"] for v in counts.values()),
        assign_fn=sum(v["fn"] for v in counts.values()),
        assign_fp=sum(v["fp"] for v in counts.values()),
        assign_misdirected=misdirected_total,
        assign_missing_not_misdirected=missing_not_misdirected_total,
        unknown_endpoints=unknown_endpoints,

        uu_gt=uu["gt"], uu_pred=uu["pred"], uu_fn=uu["fn"], uu_fp=uu["fp"], uu_mis=uu["mis"],
        up_gt=up["gt"], up_pred=up["pred"], up_fn=up["fn"], up_fp=up["fp"], up_mis=up["mis"],
        oo_gt=oo["gt"], oo_pred=oo["pred"], oo_fn=oo["fn"], oo_fp=oo["fp"], oo_mis=oo["mis"],
        op_gt=op["gt"], op_pred=op["pred"], op_fn=op["fn"], op_fp=op["fp"], op_mis=op["mis"],
    )
    return row, misdirected_edges_counter


def analyze_raw_pred_dir(input_dir: str, paths: List[str], gt_root: str) -> Tuple[List[SampleRow], Counter]:
    gt_dir = infer_gt_dir(input_dir, gt_root)
    if not os.path.isdir(gt_dir):
        raise SystemExit(f"Missing GT directory for raw-pred mode: {gt_dir}")

    gt_files = sorted(glob.glob(os.path.join(gt_dir, "*.json")))
    gt_idx = {canon_key(p, False): p for p in gt_files}
    pr_idx = {canon_key(p, True): p for p in paths}

    common = sorted(set(gt_idx.keys()) & set(pr_idx.keys()))
    missing_in_pred = sorted(set(gt_idx.keys()) - set(pr_idx.keys()))
    extra_in_pred = sorted(set(pr_idx.keys()) - set(gt_idx.keys()))

    subset = os.path.basename(gt_dir)
    print(f"[{subset}] matched={len(common)} missing_in_pred={len(missing_in_pred)} extra_in_pred={len(extra_in_pred)}")

    if not common:
        raise SystemExit(
            f"No canonical matches found.\nGT dir: {gt_dir} ({len(gt_files)} files)\nPred dir: {input_dir} ({len(paths)} files)"
        )

    rows: List[SampleRow] = []
    mis_ctr_all = Counter()
    for key in common:
        row, mis_ctr = analyze_raw_pair(gt_idx[key], pr_idx[key])
        rows.append(row)
        mis_ctr_all.update(mis_ctr)

    return rows, mis_ctr_all


# -----------------------
# Output assembly
# -----------------------

def make_markdown_table(summary: dict) -> str:
    lines = []
    lines.append("### Assignment Relation Misdirection (All Tracked Type-Pairs)\n")
    lines.append("| Category | GT (TP+FN) | Pred (TP+FP) | FN | FP | Misdirected (FN with reverse FP) | Misdirection rate (mis/FN) |")
    lines.append("|---|---:|---:|---:|---:|---:|---:|")

    for cat in [
        "OVERALL",
        "user_attributes -> user_attributes",
        "user_attributes -> policy_classes",
        "object_attributes -> object_attributes",
        "object_attributes -> policy_classes",
    ]:
        s = summary["by_category"][cat]
        rate = safe_div(s["mis"], s["fn"])
        lines.append(
            f"| {cat} | {s['gt']} | {s['pred']} | {s['fn']} | {s['fp']} | {s['mis']} | {fmt(rate)} |"
        )

    lines.append("\n**Top reversed assignment edges (by frequency)**")
    if summary["top_misdirected_edges"]:
        for item in summary["top_misdirected_edges"]:
            lines.append(f"- `{item['src']}` → `{item['dst']}`: **{item['count']}**")
    else:
        lines.append("- _(none detected)_")

    return "\n".join(lines) + "\n"


def summarize_rows(rows: List[SampleRow], top_ctr: Counter) -> dict:
    def agg(gt: int, pred: int, fn: int, fp: int, mis: int) -> dict:
        return {"gt": gt, "pred": pred, "fn": fn, "fp": fp, "mis": mis}

    summary = {
        "n_samples": len(rows),
        "unknown_endpoints_total": sum(r.unknown_endpoints for r in rows),
        "by_category": {
            "OVERALL": agg(
                sum(r.assign_gt for r in rows),
                sum(r.assign_pred for r in rows),
                sum(r.assign_fn for r in rows),
                sum(r.assign_fp for r in rows),
                sum(r.assign_misdirected for r in rows),
            ),
            "user_attributes -> user_attributes": agg(
                sum(r.uu_gt for r in rows),
                sum(r.uu_pred for r in rows),
                sum(r.uu_fn for r in rows),
                sum(r.uu_fp for r in rows),
                sum(r.uu_mis for r in rows),
            ),
            "user_attributes -> policy_classes": agg(
                sum(r.up_gt for r in rows),
                sum(r.up_pred for r in rows),
                sum(r.up_fn for r in rows),
                sum(r.up_fp for r in rows),
                sum(r.up_mis for r in rows),
            ),
            "object_attributes -> object_attributes": agg(
                sum(r.oo_gt for r in rows),
                sum(r.oo_pred for r in rows),
                sum(r.oo_fn for r in rows),
                sum(r.oo_fp for r in rows),
                sum(r.oo_mis for r in rows),
            ),
            "object_attributes -> policy_classes": agg(
                sum(r.op_gt for r in rows),
                sum(r.op_pred for r in rows),
                sum(r.op_fn for r in rows),
                sum(r.op_fp for r in rows),
                sum(r.op_mis for r in rows),
            ),
        },
        "top_misdirected_edges": [
            {"src": src, "dst": dst, "count": cnt}
            for (src, dst), cnt in top_ctr.most_common(10)
        ],
        "definition": {
            "tracked_pairs": [pair_key(a, b) for (a, b) in sorted(TRACKED_PAIRS)],
            "misdirection_rule": "FN(A->B,'assign') AND reverse(B->A,'assign') predicted",
        },
    }
    return summary


# -----------------------
# Main
# -----------------------

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True, help="Input directory containing JSON files")
    ap.add_argument("--glob", default="*.json", help="Glob for JSON files inside --dir")
    ap.add_argument("--out_prefix", default="assign_misdirection", help="Output file prefix")
    ap.add_argument("--out_dir", default=None, help="Directory to write outputs; defaults to --dir")
    ap.add_argument(
        "--gt_root",
        default="/Users/sherifdeenlawal/Venv/NIST_Projects/dataset/BlueSkyEvaluation/gt",
        help="GT root used only in raw-pred mode",
    )
    args = ap.parse_args()

    input_dir = args.dir
    out_dir = args.out_dir or input_dir
    os.makedirs(out_dir, exist_ok=True)

    paths = sorted(glob.glob(os.path.join(input_dir, args.glob)))
    if not paths:
        raise SystemExit(f"No JSON files found in {input_dir} matching {args.glob}")

    first = safe_load_json(paths[0])
    if first is None:
        raise SystemExit(f"Could not read first JSON: {paths[0]}")

    if is_post_eval_json(first):
        mode = "post_eval"
        rows: List[SampleRow] = []
        mis_ctr_all = Counter()
        for p in paths:
            row, mis_ctr = analyze_post_eval_file(p)
            rows.append(row)
            mis_ctr_all.update(mis_ctr)
    elif is_raw_pred_json(first):
        mode = "raw_pred"
        rows, mis_ctr_all = analyze_raw_pred_dir(input_dir, paths, args.gt_root)
    else:
        raise SystemExit(
            "Unrecognized JSON schema. Expected either:\n"
            "1) evaluation.relation_metrics.{fp_relations,fn_relations,...}\n"
            "or\n"
            "2) raw prediction files with nodes[] and edges[]"
        )

    summary = summarize_rows(rows, mis_ctr_all)
    summary["mode"] = mode
    summary["input_dir"] = input_dir

    out_csv = os.path.join(out_dir, f"{args.out_prefix}_per_sample.csv")
    out_json = os.path.join(out_dir, f"{args.out_prefix}_summary.json")
    out_md = os.path.join(out_dir, f"{args.out_prefix}_table.md")

    with open(out_csv, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(asdict(rows[0]).keys()))
        w.writeheader()
        for r in rows:
            w.writerow(asdict(r))

    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    md = make_markdown_table(summary)
    with open(out_md, "w", encoding="utf-8") as f:
        f.write(md)

    print(f"Mode: {mode}")
    print(md)
    print(f"Wrote: {out_csv}")
    print(f"Wrote: {out_json}")
    print(f"Wrote: {out_md}")


if __name__ == "__main__":
    main()