"""CART regressor on five players from the Hitters dataset (ISLR).

Two predictors, both numeric:
    years - seasons played in the major leagues
    hits  - hits made in the previous season
and the target:
    salary - 1987 annual salary, in thousands of dollars

Rows 1 and 2 are a genuine collision: same years, same hits, different salary.
Structure is identical to heart_tree.py; only two things change — the impurity
(variance instead of Gini) and the leaf statistic (mean instead of counts).

Run:  python3 hitters_tree.py
"""

training_data = [
    [2, 41, 67.5],    # BillyJo Robidoux
    [2, 41, 95.0],    # Jack Howell
    [3, 130, 480.0],  # Alvin Davis
    [6, 77, 670.0],   # Mike Marshall
    [7, 149, 787.5],  # Lloyd Moseby
]

names = ["BillyJo Robidoux", "Jack Howell", "Alvin Davis", "Mike Marshall", "Lloyd Moseby"]
header = ["years", "hits", "salary"]


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


# ------------------------------------------------------------- the algorithm
# Everything below is heart_tree.py with gini -> variance and counts -> mean.

def is_numeric(value):
    return isinstance(value, int) or isinstance(value, float)


def mean(rows):
    return sum(row[-1] for row in rows) / float(len(rows))


def variance(rows):
    targets = [row[-1] for row in rows]
    m = sum(targets) / len(targets)
    return sum((t - m) ** 2 for t in targets) / len(targets)


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

    def match(self, example):
        val = example[self.column]
        if is_numeric(self.value):
            return val >= self.value
        return val == self.value

    def __repr__(self):
        condition = ">=" if is_numeric(self.value) else "=="
        return "Is %s %s %s?" % (header[self.column], condition, str(self.value))


def partition(rows, question):
    true_rows, false_rows = [], []
    for row in rows:
        if question.match(row):
            true_rows.append(row)
        else:
            false_rows.append(row)
    return true_rows, false_rows


def info_gain(left, right, current_uncertainty):
    p = float(len(left)) / (len(left) + len(right))
    return current_uncertainty - p * variance(left) - (1 - p) * variance(right)


def find_best_split(rows):
    best_gain = 0
    best_question = None
    current_uncertainty = variance(rows)
    n_features = len(rows[0]) - 1

    for col in range(n_features):
        values = sorted(set(row[col] for row in rows), reverse=True)
        for val in values:
            question = Question(col, val)
            true_rows, false_rows = partition(rows, question)
            if len(true_rows) == 0 or len(false_rows) == 0:
                continue
            gain = info_gain(true_rows, false_rows, current_uncertainty)
            if gain >= best_gain:  # note the >=
                best_gain, best_question = gain, question

    return best_gain, best_question


def all_candidates(rows):
    out = []
    current = variance(rows)
    for col in range(len(rows[0]) - 1):
        for val in sorted(set(row[col] for row in rows)):
            q = Question(col, val)
            t, f = partition(rows, q)
            if not t or not f:
                out.append((None, q, len(t), len(f)))
                continue
            out.append((info_gain(t, f, current), q, len(t), len(f)))
    return out


class RegressionLeaf:
    def __init__(self, rows):
        self.prediction = mean(rows)
        self.rows = rows


class Decision_Node:
    def __init__(self, question, true_branch, false_branch):
        self.question = question
        self.true_branch = true_branch
        self.false_branch = false_branch


def build_tree(rows):
    if not rows:
        raise ValueError("Training rows must not be empty")
    gain, question = find_best_split(rows)
    if gain <= 1e-12 or question is None:
        return RegressionLeaf(rows)
    true_rows, false_rows = partition(rows, question)
    return Decision_Node(question, build_tree(true_rows), build_tree(false_rows))


def print_tree(node, spacing=""):
    if isinstance(node, RegressionLeaf):
        who = ", ".join(names[training_data.index(r)] for r in node.rows)
        print(spacing + "Predict %.2f   (%s)" % (node.prediction, who))
        return
    print(spacing + str(node.question))
    print(spacing + "--> True:")
    print_tree(node.true_branch, spacing + "  ")
    print(spacing + "--> False:")
    print_tree(node.false_branch, spacing + "  ")


def predict(row, node):
    if isinstance(node, RegressionLeaf):
        return node.prediction
    if node.question.match(row):
        return predict(row, node.true_branch)
    return predict(row, node.false_branch)


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

def mermaid(node):
    lines = ["flowchart TD"]
    leaves = []
    counter = [0]

    def walk(n):
        me = "n%d" % counter[0]
        counter[0] += 1
        if isinstance(n, RegressionLeaf):
            if len(n.rows) == 1:
                label = "%.1f" % n.prediction
            else:
                vals = " + ".join("%.1f" % r[-1] for r in n.rows)
                label = "mean(%s) = %.2f" % (vals, n.prediction)
            lines.append('    %s["%s"]' % (me, label))
            leaves.append((me, len(n.rows) > 1))
            return me
        lines.append('    %s{"%s"}' % (me, str(n.question)))
        t = walk(n.true_branch)
        f = walk(n.false_branch)
        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("%-20s %-6s %-6s %s" % ("player", header[0], header[1], header[2]))
    for nm, r in zip(names, training_data):
        print("%-20s %-6d %-6d %.1f" % (nm, r[0], r[1], r[2]))
    print("\nmean salary   : %.1f" % mean(training_data))
    print("root variance : %.2f" % variance(training_data))

    section("2. every candidate split at the root")
    rows = []
    for gain, q, nt, nf in all_candidates(training_data):
        label = "     -- does not divide" if gain is None else "%18.10f" % gain
        rows.append((label, q, nt, nf))
    for g, q, nt, nf in sorted(rows, key=lambda r: r[0], reverse=True):
        print("  %-20s %-24s  true=%d false=%d" % (g, q, nt, nf))
    best_gain, best_q = find_best_split(training_data)
    print("\n  winner: %s   gain=%r" % (best_q, best_gain))

    section("3. the fitted tree")
    tree = build_tree(training_data)
    print_tree(tree)

    section("4. predictions on the training rows")
    for nm, r in zip(names, training_data):
        print("  %-20s actual %7.1f   predicted %7.2f" % (nm, r[-1], predict(r, tree)))

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