"""
Basic example of using driedfigs to capture figure data.

This example demonstrates the core workflow:
1. Start capturing data for a figure
2. Create plots with the 'name' parameter
3. Save both the figure and its data
"""

import driedfigs as df
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

# Create output directories
output_base = Path("example_output")
figure_dir = output_base / "figures"
data_dir = output_base / "data"

figure_dir.mkdir(parents=True, exist_ok=True)
data_dir.mkdir(parents=True, exist_ok=True)

# Generate some example data
x = np.linspace(0, 2 * np.pi, 100)
y_sin = np.sin(x)
y_cos = np.cos(x)

# Add some noise to make it more realistic
noise = np.random.normal(0, 0.1, 100)
y_measurements = y_sin + noise

# Start capturing data for this figure
print("Starting figure data capture...")
fig_data = df.start_figure("trigonometric_example", verbose=True)

# Create the plot
fig, ax = plt.subplots(figsize=(10, 6))

# Plot with 'name' parameter to capture data
ax.plot(x, y_sin, '-', label="sin(x)", linewidth=2, name="sine_wave")
ax.plot(x, y_cos, '--', label="cos(x)", linewidth=2, name="cosine_wave")
ax.scatter(x[::10], y_measurements[::10],
           label="Measurements", alpha=0.6, name="measurement_points")

ax.set_xlabel("x (radians)")
ax.set_ylabel("y")
ax.set_title("Trigonometric Functions with Measurements")
ax.legend()
ax.grid(True, alpha=0.3)

# Save the data
print("\nSaving data...")
df.save_figuredata(data_dir / "trigonometric_example.pkl")

# Save the figure manually
print("Saving figure...")
fig.savefig(figure_dir / "trigonometric_example.pdf", bbox_inches='tight')
fig.savefig(figure_dir / "trigonometric_example.png", bbox_inches='tight', dpi=150)

print(f"\n✓ Figure saved to: {figure_dir}")
print(f"✓ Data saved to: {data_dir}")

# Demonstrate loading the saved data
print("\nLoading saved data to verify...")
import pickle

data_file = data_dir / "trigonometric_example.pkl"
with open(data_file, "rb") as f:
    loaded_data = pickle.load(f)

print(f"\nCaptured data contains {len(loaded_data)} named plots:")
for name in loaded_data.keys():
    print(f"  - {name}")

# Access specific plot data
sine_data = loaded_data["sine_wave"]
print(f"\nSine wave data shape: {sine_data[0].shape}, {sine_data[1].shape}")
print(f"First few x values: {sine_data[0][:5]}")
print(f"First few y values: {sine_data[1][:5]}")

print("\n✓ Example completed successfully!")
print(f"  Check {figure_dir} for the figures")
print(f"  Check {data_dir} for the data pickle file")
