import os
import sys
from typing import List, Tuple

# Ensure python folder is in path
python_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, python_dir)

import torch  # noqa: E402
import fces_native  # noqa: E402
from send_telemetry import push_to_mariadb, push_to_surrealdb  # noqa: E402


def ackley(w: torch.Tensor) -> torch.Tensor:
    x, y = w[0], w[1]
    part1 = -20.0 * torch.exp(-0.2 * torch.sqrt(0.5 * (x**2 + y**2)))
    part2 = -torch.exp(
        0.5 * (torch.cos(2 * 3.1415926535 * x) + torch.cos(2 * 3.1415926535 * y))
    )
    return part1 + part2 + 2.718281828459 + 20.0


def run_adamw(steps: int = 200) -> Tuple[List[float], torch.Tensor]:
    torch.manual_seed(42)
    # Start at x=3.0, y=3.0 (trapped in a local minimum)
    w = torch.tensor([3.0, 3.0])
    w.requires_grad_(True)
    optimizer = torch.optim.AdamW([w], lr=0.1)

    losses = []
    for step in range(steps):
        optimizer.zero_grad()
        loss = ackley(w)
        loss.backward()
        optimizer.step()
        losses.append(float(loss.item()))
    return losses, w.detach().clone()


def run_fces(steps: int = 200) -> Tuple[List[float], torch.Tensor]:
    torch.manual_seed(42)
    # Start at x=3.0, y=3.0
    w = torch.tensor([3.0, 3.0])
    w.requires_grad_(True)

    cfg = fces_native.FCESConfig()
    cfg.lr = 0.1
    cfg.population_size = 64
    cfg.total_steps = steps

    optimizer = fces_native.FCESOptimizer([w], cfg)

    losses = []
    for step in range(steps):
        optimizer.zero_grad()
        loss = ackley(w)
        loss.backward()
        optimizer.step()
        optimizer.update_fitness(float(loss.item()))
        losses.append(float(loss.item()))
    return losses, w.detach().clone()


def main() -> None:
    print("=" * 80)
    print(
        "      FCES VS ADAMW CONVERGENCE BENCHMARK (NON-CONVEX ACKLEY FUNCTION)      "
    )
    print("=" * 80)

    steps = 200

    print("[INFO] Running AdamW baseline...")
    adam_losses, adam_final_w = run_adamw(steps)

    print("[INFO] Running FCES optimizer...")
    fces_losses, fces_final_w = run_fces(steps)

    print("\n" + "-" * 40 + " BENCHMARK RESULT SUMMARY " + "-" * 40)
    print(f"{'Step':<10}{'AdamW Loss':<20}{'FCES Loss':<20}{'FCES Improvement':<20}")
    print("-" * 106)

    telemetry_entries = []
    for idx in range(0, steps, 20):
        improvement = adam_losses[idx] - fces_losses[idx]
        pct = (improvement / adam_losses[idx]) * 100 if adam_losses[idx] > 0 else 0
        print(f"{idx:<10}{adam_losses[idx]:<20.6f}{fces_losses[idx]:<20.6f}{pct:.2f}%")

        telemetry_entries.append(
            (
                "INFO",
                "benchmark_step",
                f"Step {idx} | AdamW Loss: {adam_losses[idx]:.4f} | FCES Loss: {fces_losses[idx]:.4f}",
            )
        )

    print("-" * 106)
    print(f"Final AdamW Loss: {adam_losses[-1]:.6f} (Final w: {adam_final_w.numpy()})")
    print(f"Final FCES Loss:  {fces_losses[-1]:.6f} (Final w: {fces_final_w.numpy()})")

    # Calculate ratio improvement
    improvement_factor = adam_losses[-1] - fces_losses[-1]
    print(f"FCES Absolute Improvement: {improvement_factor:.6f} lower loss!")
    print("=" * 80)

    # Send telemetry
    push_to_surrealdb(telemetry_entries)
    push_to_mariadb(telemetry_entries)

    # Save results to a summary file
    summary_path = os.path.join(os.path.dirname(python_dir), "benchmark_summary.txt")
    with open(summary_path, "w") as f:
        f.write("FCES vs AdamW Optimization Benchmark (Ackley Function)\n")
        f.write(
            f"Final AdamW Loss: {adam_losses[-1]:.6f} (w: {adam_final_w.numpy()})\n"
        )
        f.write(
            f"Final FCES Loss:  {fces_losses[-1]:.6f} (w: {fces_final_w.numpy()})\n"
        )
        f.write(f"Absolute Improvement: {improvement_factor:.6f}\n")
    print(f"[INFO] Benchmark summary saved to {summary_path}")


if __name__ == "__main__":
    main()