"""
Example: Basic Operations and Autograd

This example demonstrates basic scalar operations and automatic differentiation.

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

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

from scalargrad import Scalar, console
from rich.rule import Rule
from rich.table import Table
from rich import box


def main():
    console.print(Rule("[header]Basic Operations and Autograd[/header]", style="banner_border"))

    # ── Computation graph & gradients ──────────────────────────────────────
    console.print()
    a = Scalar(-4.0, label='a')
    b = Scalar(2.0, label='b')

    # Build computation graph
    c = a + b
    c.label = 'c'
    d = a * b + b ** 3
    d.label = 'd'
    e = c - d
    e.label = 'e'
    f = e ** 2
    f.label = 'f'
    g = f / 2.0
    g.label = 'g'
    g = g + 10.0 / f

    console.print(f"[label]Final value:[/label] [value]{g}[/value]")
    console.print("[info]Computing gradients...[/info]")
    g.backward()

    grad_table = Table(box=box.SIMPLE_HEAD, border_style="table_border", header_style="table_header")
    grad_table.add_column("Variable", style="label", justify="center")
    grad_table.add_column("Gradient",  style="value",  justify="right")
    grad_table.add_row("dg/da", f"{a.grad:.4f}")
    grad_table.add_row("dg/db", f"{b.grad:.4f}")
    console.print(grad_table)

    # ── Activation functions ───────────────────────────────────────────────
    console.print(Rule("[header]Activation Functions[/header]", style="banner_border"))
    console.print()

    x = Scalar(0.5, label='x')

    act_table = Table(
        box=box.ROUNDED,
        border_style="table_border",
        header_style="table_header",
        show_lines=False,
    )
    act_table.add_column("Function",  style="label",   justify="left",  width=12)
    act_table.add_column("Input",     style="key",     justify="right", width=8)
    act_table.add_column("Output",    style="value",   justify="right", width=14)

    relu_out    = x.relu()
    tanh_out    = x.tanh()
    sigmoid_out = x.sigmoid()
    exp_out     = x.exp()
    log_out     = x.log()

    act_table.add_row("ReLU",    f"{x.data}", f"{relu_out.data:.4f}")
    act_table.add_row("tanh",    f"{x.data}", f"{tanh_out.data:.4f}")
    act_table.add_row("sigmoid", f"{x.data}", f"{sigmoid_out.data:.4f}")
    act_table.add_row("exp",     f"{x.data}", f"{exp_out.data:.4f}")
    act_table.add_row("log",     f"{x.data}", f"{log_out.data:.4f}")

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


if __name__ == "__main__":
    main()
