import torch
from math import pi
import matplotlib.pyplot as plt

import sys
sys.path.append("/home/judson/Neural-Networks-in-GNC/inverted_pendulum")
from training.base_loss_functions import base_loss_functions

# Create an array of error values from 0 to 2π
errors = torch.linspace(0, 2 * pi, 1000)
desired = torch.zeros_like(errors)  # Assume desired_theta = 0 for plotting

plt.figure(figsize=(10, 6))
for name, (exponent, loss_fn) in base_loss_functions.items():
    # Compute loss for each error value
    losses = loss_fn(errors, desired, min_val=0.01)
    plt.plot(errors.numpy(), losses.numpy(), label=f"p={exponent:.2f}")

plt.xlabel(r"Error ($|\theta - \theta_d|$)")
plt.ylabel("Loss")
plt.title("Base Loss Functions")
plt.legend()
plt.grid(True)
plt.savefig("base_loss_functions.png")