"""
Example: Neural Network Training

This example demonstrates training a neural network on synthetic data.

To change the output theme, set config.theme before running::

    from scalargrad import config
    config.theme = "dark"     # or: cyberpunk / light / minimal
"""

import warnings
warnings.filterwarnings('ignore')

from scalargrad import MLP, Adam, SVMLoss, Trainer, Visualizer, config, LogLevel, console
from rich.rule import Rule
from rich.table import Table
from rich import box

# Check for optional dependencies
try:
    from sklearn.datasets import make_moons
    import numpy as np
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False
    console.print("[warning]scikit-learn not available.[/warning] Install with: [code]pip install scikit-learn[/code]")

try:
    import matplotlib.pyplot as plt
    PLOTTING_AVAILABLE = True
except ImportError:
    PLOTTING_AVAILABLE = False
    console.print("[warning]Matplotlib not available.[/warning] Install with: [code]pip install matplotlib[/code]")

try:
    import graphviz as _gv  # noqa: F401
    GRAPHVIZ_AVAILABLE = True
except ImportError:
    GRAPHVIZ_AVAILABLE = False


def main():
    if not SKLEARN_AVAILABLE:
        console.print("[error]This example requires scikit-learn.[/error] Install with: [code]pip install scikit-learn[/code]")
        return

    console.print(Rule("[header]Neural Network Training[/header]", style="banner_border"))

    # Configure
    config.log_level = LogLevel.INFO
    config.seed = 42

    # ── Dataset ─────────────────────────────────────────────────────────────
    X, y = make_moons(n_samples=100, noise=0.1, random_state=42)
    y = y * 2 - 1  # Convert to {-1, 1}

    dataset_table = Table(box=box.SIMPLE_HEAD, border_style="table_border", header_style="table_header", show_header=False)
    dataset_table.add_column("Key",   style="label")
    dataset_table.add_column("Value", style="value")
    dataset_table.add_row("Samples",  str(X.shape[0]))
    dataset_table.add_row("Features", str(X.shape[1]))
    dataset_table.add_row("Classes",  str(np.unique(y).tolist()))
    console.print("\n[subheader]Dataset[/subheader]")
    console.print(dataset_table)

    # ── Model ────────────────────────────────────────────────────────────────
    model = MLP(
        nin=2,
        layer_sizes=[16, 16, 1],
        activations=['relu', 'relu', 'linear']
    )

    console.print("\n[subheader]Model[/subheader]")
    console.print(model)   # uses __rich_repr__
    console.print(f"  [label]Parameters:[/label] [value]{model.num_parameters()}[/value]")

    # ── Training ─────────────────────────────────────────────────────────────
    optimizer = Adam(model.parameters(), lr=0.01)
    loss_fn = SVMLoss()
    trainer = Trainer(model, optimizer, loss_fn)

    console.print()
    history = trainer.fit(X.tolist(), y.tolist(), epochs=50, verbose=True)

    # ── Final results ─────────────────────────────────────────────────────────
    final_metrics = history[-1]
    results_table = Table(
        title="[table_header]Final Results[/table_header]",
        box=box.ROUNDED,
        border_style="table_border",
        header_style="table_header",
    )
    results_table.add_column("Metric",    style="label")
    results_table.add_column("Value",     style="value",   justify="right")
    results_table.add_row("Loss",         f"{final_metrics.loss:.4f}")
    results_table.add_row("Accuracy",     f"{final_metrics.accuracy:.2%}")
    console.print(results_table)

    # ── Visualization ─────────────────────────────────────────────────────────
    if PLOTTING_AVAILABLE or GRAPHVIZ_AVAILABLE:
        console.print("\n[info]Generating visualizations...[/info]")

    if PLOTTING_AVAILABLE:
        Visualizer.plot_training_history(history)
    else:
        console.print("[warning]Skipping training history plot[/warning] (matplotlib not available)")

    # ── Computation graph of a single forward pass ────────────────────────
    if GRAPHVIZ_AVAILABLE:
        console.print("[info]Rendering computation graph for one sample...[/info]")
        try:
            from graphviz.backend.execute import ExecutableNotFound
            from scalargrad import Scalar
            sample = [Scalar(float(v)) for v in X[0]]
            out = model(sample)
            dot = Visualizer.draw_graph(out, format="svg")
            dot.render("computation_graph", cleanup=True)
            console.print(
                "[success]>> Computation graph saved:[/success] "
                "[value]computation_graph.svg[/value]"
            )
        except ExecutableNotFound:
            console.print(
                "[warning]Skipping computation graph[/warning] — Graphviz system binary not found.\n"
                "  Install it from [value]https://graphviz.org/download/[/value] and add it to PATH,\n"
                "  or on Windows run: [code]winget install graphviz[/code]"
            )
    else:
        console.print("[warning]Skipping computation graph[/warning] (graphviz Python package not available — "
                      "install with: [code]pip install graphviz[/code])")

    if PLOTTING_AVAILABLE:
        Visualizer.plot_decision_boundary(model, X, y)
    elif not PLOTTING_AVAILABLE:
        console.print("[warning]Skipping decision boundary plot[/warning] (matplotlib not available)")

    console.print(Rule(style="banner_border"))


if __name__ == "__main__":
    main()