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.time_weighting_functions import weight_functions

# Time grid
t_start, t_end, t_points = 0, 10, 1000
t_span = torch.linspace(t_start, t_end, t_points)

plt.figure(figsize=(10, 6))

# Dictionaries to separate non-mirrored and mirrored functions and to store colors
colors = {}

# First plot non-mirrored functions and record their colors
for name, weight_fn in weight_functions.items():
    if "mirrored" not in name:
        weights = weight_fn(t_span, t_max=t_end, min_val=0.01)
        line, = plt.plot(t_span.numpy(), weights.numpy(), label=name, linestyle="-")
        colors[name] = line.get_color()

# Now plot mirrored functions with the same color as their non-mirrored counterpart
for name, weight_fn in weight_functions.items():
    if "mirrored" in name:
        weights = weight_fn(t_span, t_max=t_end, min_val=0.01)
        # Derive base name; assuming the mirrored function is named as "base_mirrored"
        base_name = name.replace("_mirrored", "")
        color = colors.get(base_name, None)
        plt.plot(t_span.numpy(), weights.numpy(), label=name, linestyle="-.", color=color)

plt.axvline(x=5, label="Mirror Line", color='black', linestyle='--')
plt.xlabel(r"Time of Simulation ($t_i$)")
plt.ylabel("Weight")
plt.title("Time Weighting Functions and Their Mirror")
plt.legend()
plt.grid(True)
plt.savefig("time_weighting_functions.png")