"""Train with CoolBox, or render saved native predictions with --results FILE."""

import argparse
import json
from html import escape
from math import log
from pathlib import Path


X_train = [
    [0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0],
    [0.1, 0.1], [0.1, 0.9], [0.9, 0.1], [0.9, 0.9],
]
y_train = [0, 1, 1, 0, 0, 1, 1, 0]
X_validation = [
    [0.25, 0.25], [0.25, 0.75], [0.75, 0.25], [0.75, 0.75],
    [0.35, 0.35], [0.35, 0.65], [0.65, 0.35], [0.65, 0.65],
]
y_validation = [0, 1, 1, 0, 0, 1, 1, 0]
X_test = [[0.05, 0.8], [0.8, 0.05], [0.05, 0.05], [0.8, 0.8]]
y_test = [1, 1, 0, 0]
GRID_SIZE = 60


def log_loss(model, rows, labels):
    class_index = {label: index for index, label in enumerate(model.classes())}
    return -sum(
        log(max(1e-15, model.predict_proba(row)[class_index[label]]))
        for row, label in zip(rows, labels)
    ) / len(labels)


def train():
    from ml_toolbox import decision_tree as trees, opt

    def decode(genes):
        params = trees.BoostTreeParameters()
        params.num_estimators = 20 + int(80 * genes[0] + 0.5)
        params.learning_rate = 0.03 + 0.17 * genes[1]
        params.max_depth = 1 + int(3 * genes[2] + 0.5)
        return params

    def fit(genes):
        model = trees.BoostTreeClassifier(decode(genes))
        model.fit(X_train, y_train)
        return model

    initial = [0.5, (0.1 - 0.03) / 0.17, 0.0]  # 60 trees, rate 0.1, depth 1
    baseline = fit(initial)
    baseline_loss = log_loss(baseline, X_validation, y_validation)
    # dimensions, population, generations, tournament, crossover, mutation,
    # mutation sigma, lower bound, upper bound, seed (native declaration order).
    search = opt.Optimizer(
        opt.OptimizerType.GENETIC_SEARCH,
        [3, 12, 8, 3, 0.85, 0.2, 0.15, 0.0, 1.0, 42],
    )
    winner = search.optimize(
        lambda genes: log_loss(fit(genes), X_validation, y_validation), initial
    )
    candidate = fit(winner)
    improved = log_loss(candidate, X_validation, y_validation) < baseline_loss
    selected = candidate if improved else baseline
    selected_genes = winner if improved else initial
    models = []
    for model, genes in [(baseline, initial), (selected, selected_genes)]:
        params = decode(genes)
        class_one = list(model.classes()).index(1)
        cells = []
        for row in range(GRID_SIZE):
            for column in range(GRID_SIZE):
                point = [(column + 0.5) / GRID_SIZE, 1 - (row + 0.5) / GRID_SIZE]
                cells.append([model.predict(point), model.predict_proba(point)[class_one]])
        models.append({
            "num_estimators": params.num_estimators,
            "learning_rate": params.learning_rate,
            "max_depth": params.max_depth,
            "validation_loss": log_loss(model, X_validation, y_validation),
            "test_accuracy": sum(
                model.predict(x) == y for x, y in zip(X_test, y_test)
            ) / len(y_test),
            "cells": cells,
        })
    return {"grid_size": GRID_SIZE, "models": models}


def text(x, y, value, size=12, anchor="start", color="#596579", extra=""):
    return (
        f'<text x="{x}" y="{y}" font-family="Arial, sans-serif" font-size="{size}" '
        f'text-anchor="{anchor}" fill="{color}" {extra}>{escape(value)}</text>'
    )


def panel(model, x, y, width, height, grid):
    parts = []
    for index, (label, probability) in enumerate(model["cells"]):
        row, column = divmod(index, grid)
        color = "#435da8" if label == 0 else "#a45b29"
        opacity = 0.12 + 0.7 * abs(2 * probability - 1)
        parts.append(
            f'<rect x="{x + column * width / grid:.3f}" '
            f'y="{y + row * height / grid:.3f}" width="{width / grid:.3f}" '
            f'height="{height / grid:.3f}" fill="{color}" fill-opacity="{opacity:.4f}"/>'
        )
    parts.append(
        f'<path d="M{x} {y}V{y + height}H{x + width}" fill="none" stroke="#9aa5b5"/>'
    )
    for value in (0, 0.5, 1):
        parts.append(text(x + value * width, y + height + 23, f"{value:g}", anchor="middle"))
        parts.append(text(x - 14, y + height * (1 - value) + 4, f"{value:g}", anchor="end"))
    parts.append(text(x + width / 2, y + height + 48, "Feature 1", anchor="middle"))
    parts.append(text(
        x - 48, y + height / 2, "Feature 2", anchor="middle",
        extra=f'transform="rotate(-90 {x - 48} {y + height / 2})"',
    ))
    for point, label in zip(X_train, y_train):
        color = "#435da8" if label == 0 else "#a45b29"
        parts.append(
            f'<circle cx="{x + point[0] * width}" cy="{y + (1 - point[1]) * height}" '
            f'r="5" fill="{color}" stroke="white" stroke-width="2"/>'
        )
    for point in X_test:
        cx, cy = x + point[0] * width, y + (1 - point[1]) * height
        parts.append(
            f'<path d="M{cx} {cy - 6}l6 6-6 6-6-6z" fill="white" '
            'stroke="#263246" stroke-width="1.5"/>'
        )
    return parts


def write_surface(results, destination, comparison=True):
    models, grid = results["models"], results["grid_size"]
    if len(models) != 2 or grid != GRID_SIZE:
        raise ValueError("Expected two models on a 60-by-60 grid")
    for model in models:
        if len(model["cells"]) != grid * grid:
            raise ValueError("Incomplete prediction grid")
        for label, probability in model["cells"]:
            if label not in (0, 1) or not 0 <= probability <= 1:
                raise ValueError("Invalid class or probability in prediction grid")
    title = "Genetic search: measured comparison" if comparison else "Boosted trees: shallow baseline"
    parts = [
        '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 960 560" '
        'role="img" aria-labelledby="title description">',
        f'<title id="title">{escape(title)}</title>',
        '<desc id="description">Actual CoolBox predictions. Blue is class zero, orange '
        'is class one. The same shade scale represents confidence in each plot. '
        'Circles show training labels; diamonds mark held-out test locations.</desc>',
        '<rect width="960" height="560" fill="white"/>',
        text(54, 40, title, 22, color="#263246"),
        text(54, 65, "Same data, axes, markers and confidence scale; no illustrative regions."),
    ]
    layout = [(models[0], 90, "Before: baseline"), (models[1], 550, "After: selected")]
    if not comparison:
        layout = [(models[0], 132, "60 trees, rate 0.1, depth 1")]
    for model, offset, heading in layout:
        if comparison:
            params = (
                f'{model["num_estimators"]} trees | rate {model["learning_rate"]:.4f} '
                f'| depth {model["max_depth"]}'
            )
            parts += [text(offset, 98, heading, 17, color="#263246"), text(offset, 119, params)]
            y, width, height = 166, 300, 260
            metrics_y = 141
        else:
            parts.append(text(offset, 94, heading, 15, color="#263246"))
            y, width, height, metrics_y = 138, 660, 288, 117
        parts.append(text(
            offset, metrics_y,
            f'Validation loss {model["validation_loss"]:.4f} | test {model["test_accuracy"]:.0%}',
        ))
        parts.extend(panel(model, offset, y, width, height, grid))
    parts += [
        text(54, 503, "Blue: class 0 | Orange: class 1 | Darker: more confident | Circles: training | Diamonds: test"),
        text(54, 534, "Measured on this tiny teaching dataset; not a generalization benchmark.", 11),
        "</svg>",
    ]
    destination.write_text("\n".join(parts), encoding="utf-8")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results", type=Path, help="Render saved native predictions without importing CoolBox")
    parser.add_argument("--output-dir", type=Path, default=Path.cwd())
    args = parser.parse_args()
    results = json.loads(args.results.read_text(encoding="utf-8")) if args.results else train()
    if not args.output_dir.is_dir():
        parser.error("--output-dir must be an existing directory")
    if not args.results:
        (args.output_dir / "boosted-tree-search-results.json").write_text(
            json.dumps(results), encoding="utf-8"
        )
    write_surface(results, args.output_dir / "boosted-tree-search-preview.svg")
    write_surface(results, args.output_dir / "boosted-tree-decision-surface.svg", comparison=False)
    for name, model in zip(("Baseline", "Selected"), results["models"]):
        print(f'{name}: depth {model["max_depth"]}, validation loss '
              f'{model["validation_loss"]:.6f}, test accuracy {model["test_accuracy"]:.0%}')
    print(f"Saved matching surfaces in {args.output_dir.resolve()}")


if __name__ == "__main__":
    main()