"""CART classifier on five patients from the Cleveland heart-disease study.

Two predictors kept out of the original thirteen:
    stress_test - thallium stress test: normal / fixed defect / reversable defect
    vessels     - number of major vessels (0-3) colored by fluoroscopy
and the outcome:
    disease     - did an angiogram find heart disease?

Rows 2 and 5 are a genuine collision: same stress_test, same vessels,
opposite outcome.
Algorithm is the one the article builds, unchanged.

Run:  python3 heart_tree.py
"""

training_data = [
    ["normal", 0, "No"],
    ["fixed", 0, "Yes"],
    ["reversable", 2, "Yes"],
    ["reversable", 1, "Yes"],
    ["fixed", 0, "No"],
]

FEATURES = ["stress_test", "vessels", "disease"]


def section(title):
    print()
    print("=" * 66)
    print(title)
    print("=" * 66)


# ------------------------------------------------------------- the algorithm

def label_counts(rows):
    """Tally the labels in a pile — the label is always the last column."""
    counts = {}
    for row in rows:
        counts[row[-1]] = counts.get(row[-1], 0) + 1
    return counts


def gini(rows):
    """Impurity of a pile: 0 when every row in it carries the same label."""
    if not rows:
        raise ValueError("Cannot measure impurity of an empty group")
    impurity = 1
    for count in label_counts(rows).values():
        share = count / len(rows)
        impurity -= share ** 2
    return impurity


class Rule:
    """One yes/no test: a column, and the value it is compared against."""

    def __init__(self, column, value):
        self.column = column
        self.value = value

    def holds(self, row):
        observed = row[self.column]
        if isinstance(self.value, (int, float)):
            return observed >= self.value     # numeric: threshold
        return observed == self.value         # categorical: equality

    def __repr__(self):
        operator = ">=" if isinstance(self.value, (int, float)) else "=="
        return f"Is {FEATURES[self.column]} {operator} {self.value}?"


def split_rows(rows, rule):
    """Sort every row into the pile where the rule holds, and the pile where it does not."""
    true_pile, false_pile = [], []
    for row in rows:
        (true_pile if rule.holds(row) else false_pile).append(row)
    return true_pile, false_pile


def split_gain(parent_impurity, true_pile, false_pile):
    """What went in, minus the two piles that came out, each weighed by its share."""
    share = len(true_pile) / (len(true_pile) + len(false_pile))
    return parent_impurity - share * gini(true_pile) - (1 - share) * gini(false_pile)


def choose_split(rows):
    parent_impurity = gini(rows)
    winning_gain, winning_rule = 0, None

    for column in range(len(rows[0]) - 1):
        for value in sorted(set(row[column] for row in rows), reverse=True):
            rule = Rule(column, value)
            true_pile, false_pile = split_rows(rows, rule)

            if not true_pile or not false_pile:
                continue  # this rule doesn't divide the data

            gain = split_gain(parent_impurity, true_pile, false_pile)

            if gain >= winning_gain:  # note the >=
                winning_gain, winning_rule = gain, rule

    return winning_gain, winning_rule


def every_candidate(rows):
    """Every candidate rule and its gain, for inspection."""
    out = []
    parent_impurity = gini(rows)
    for column in range(len(rows[0]) - 1):
        for value in sorted(set(row[column] for row in rows), key=str):
            rule = Rule(column, value)
            true_pile, false_pile = split_rows(rows, rule)
            if not true_pile or not false_pile:
                out.append((None, rule, len(true_pile), len(false_pile)))
                continue
            gain = split_gain(parent_impurity, true_pile, false_pile)
            out.append((gain, rule, len(true_pile), len(false_pile)))
    return out


class Leaf:
    def __init__(self, rows):
        self.counts = label_counts(rows)
        self.rows = rows


class Node:
    def __init__(self, rule, if_true, if_false):
        self.rule = rule
        self.if_true = if_true
        self.if_false = if_false


def grow_tree(rows):
    if not rows:
        raise ValueError("Training rows must not be empty")
    gain, rule = choose_split(rows)
    if gain <= 1e-12 or rule is None:
        return Leaf(rows)
    true_pile, false_pile = split_rows(rows, rule)
    return Node(rule, grow_tree(true_pile), grow_tree(false_pile))


def show_tree(node, indent=""):
    if isinstance(node, Leaf):
        print(indent + "Predict", node.counts)
        return
    print(indent + str(node.rule))
    print(indent + "--> True:")
    show_tree(node.if_true, indent + "  ")
    print(indent + "--> False:")
    show_tree(node.if_false, indent + "  ")


def descend(row, node):
    if isinstance(node, Leaf):
        return node.counts
    branch = node.if_true if node.rule.holds(row) else node.if_false
    return descend(row, branch)


def as_percentages(counts):
    total = sum(counts.values())
    return {label: f"{count / total:.0%}" for label, count in counts.items()}


# ------------------------------------------------------------------- mermaid

def mermaid(node):
    """Emit the fitted tree as a mermaid flowchart."""
    lines = ["flowchart TD"]
    leaves = []
    counter = [0]

    def walk(n):
        me = "n%d" % counter[0]
        counter[0] += 1
        if isinstance(n, Leaf):
            total = sum(n.counts.values())
            parts = ", ".join(
                "%s %d/%d" % (k, v, total) for k, v in sorted(n.counts.items())
            )
            lines.append('    %s["%s"]' % (me, parts))
            leaves.append((me, len(n.counts) > 1))
            return me
        lines.append('    %s{"%s"}' % (me, str(n.rule)))
        t = walk(n.if_true)
        f = walk(n.if_false)
        lines.append("    %s -->|True| %s" % (me, t))
        lines.append("    %s -->|False| %s" % (me, f))
        return me

    walk(node)
    pure = [i for i, mixed in leaves if not mixed]
    mixed = [i for i, m in leaves if m]
    lines.append("    classDef pure fill:#dcf5e3,stroke:#3ab54a,color:#1c1c22;")
    lines.append("    classDef mixed fill:#fde2e0,stroke:#ef3b2c,color:#1c1c22;")
    if pure:
        lines.append("    class %s pure;" % ",".join(pure))
    if mixed:
        lines.append("    class %s mixed;" % ",".join(mixed))
    return "\n".join(lines)


# ----------------------------------------------------------------------- run

if __name__ == "__main__":
    section("1. the data")
    print("%-12s %-8s %s" % tuple(FEATURES))
    for r in training_data:
        print("%-12s %-8s %s" % tuple(map(str, r)))
    print("\nlabel_counts:", label_counts(training_data))
    print("root gini   : %r" % gini(training_data))

    n = len(training_data)
    disagree = sum(1 for a in training_data for b in training_data if a[-1] != b[-1])
    print("ordered pairs that disagree: %d/%d = %s" % (disagree, n * n, disagree / (n * n)))

    section("2. every candidate split at the root")
    rows = []
    for gain, rule, n_true, n_false in every_candidate(training_data):
        if gain is None:
            rows.append(("     -- does not divide the data", rule, n_true, n_false))
        else:
            rows.append(("%.16f" % gain, rule, n_true, n_false))
    for g, rule, n_true, n_false in sorted(rows, key=lambda r: r[0], reverse=True):
        print("  %-20s %-28s  true=%d false=%d" % (g, rule, n_true, n_false))
    winning_gain, winning_rule = choose_split(training_data)
    print("\n  winner: %s   gain=%r" % (winning_rule, winning_gain))

    section("3. the fitted tree")
    tree = grow_tree(training_data)
    show_tree(tree)

    section("4. predictions on the training rows")
    for r in training_data:
        print("  %-28s -> %s" % (str(r), as_percentages(descend(r, tree))))

    section("5. a patient with a missing stress_test value")
    unseen = ["NA", 0, "?"]
    print("  %-28s -> %s" % (str(unseen), as_percentages(descend(unseen, tree))))
    print("  (no branch for it; the == test simply fails and it slides False)")

    section("6. mermaid")
    print(mermaid(tree))