from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

if __name__ == "__main__":
    # Path to your .npy file with anomaly scores
    scores_path = Path(
        "/home/fedex/mt/projects/thesis-kowalczyk-jan/Deep-SAD-PyTorch/infer/"
        "DeepSAD/subter_test/inference/3_smoke_human_walking_2023-01-23.npy"
    )

    # Frames per second
    fps = 10

    # Load all scores
    all_scores = np.load(scores_path)

    # (Optional) If your scores are a 1D array, remove or adjust this reshape
    # The final shape is expected to be (1, total_frames) or (N,) if you want only one line.
    all_scores = all_scores.reshape(1, -1)  # shape becomes (1, total_frames)

    # We only want frames [0..299] => 300 frames
    max_frames = 300
    scores = all_scores[0, :max_frames]  # Take the first row's first 300 frames

    # Convert frames to time in seconds: time[i] = i / fps
    time_values = np.arange(len(scores)) / fps  # shape is (max_frames,)

    # Recalculate y-limit based on the selected slice
    y_limit = 1.1 * np.max(scores[np.isfinite(scores)])
    # If you want to ensure y_limit >= 10, uncomment below:
    # y_limit = max(y_limit, 10)

    print("Selected Scores Shape:", scores.shape)
    print("Time Range:", time_values[0], "->", time_values[-1])
    print("Y-limit:", y_limit)

    # --------------------------
    # Improve default font sizes
    # --------------------------
    plt.rcParams.update(
        {
            "figure.dpi": 150,
            "font.size": 14,
            "axes.labelsize": 16,
            "axes.titlesize": 16,
            "xtick.labelsize": 14,
            "ytick.labelsize": 14,
            "legend.fontsize": 14,
        }
    )

    # --------------------------------------------------
    # Create the figure (half the previous width = 3.84)
    # --------------------------------------------------
    fig, ax = plt.subplots(figsize=(3.84, 7))

    # Plot the scores vs. time
    ax.plot(time_values, scores, label="Anomaly Score", color="blue")

    # Set axis labels
    ax.set_xlabel("Time [s]", fontsize=16)
    ax.set_ylabel("Score", fontsize=16)

    # Restrict X-axis from 0 to last time value (or 30 s if exactly 300 frames)
    ax.set_xlim([0, time_values[-1]])
    ax.set_ylim([0, y_limit])

    # -------------------------
    # Customize the ticks
    # -------------------------
    # Example: only 4 ticks on the x-axis (0, 10, 20, 30) if 300 frames => 30 s
    ax.set_xticks([0, 10, 20, 30])

    # Guarantee at least one tick at y=10.
    # Adjust other ticks as you like; here we use 0 and the y_limit.
    # If y_limit <= 10, you might want to override it.
    if y_limit > 10:
        ax.set_yticks([0, 10, round(y_limit, 1)])  # round for cleaner display
    else:
        # If your data doesn't go that high, just keep 0 and y_limit
        ax.set_yticks([0, round(y_limit, 1)])

    # Optionally add a legend
    ax.legend(loc="upper right")

    plt.axvline(x=22, color="r", label="shown image")

    # Tight layout for small figure
    plt.tight_layout()

    # ----------------------
    # Save to disk, no show
    # ----------------------
    output_filename = f"{scores_path.stem}_static_time_plot.png"
    plt.savefig(output_filename, dpi=150)

    # Clear the figure if you continue using matplotlib
    plt.clf()