import os
import xarray as xr
#import matplotlib
#matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# ============================================================
# Figure settings
# ============================================================
FDIR = "../out/"

script_name = os.path.splitext(os.path.basename(__file__))[0]

outfile_mp4 = script_name + ".mp4"
outfile_png = script_name + ".png"

TSTR = 0
TEND = 20

FRAME_SKIP = 100
HOLD_FRAMES = 10      # hold the final view

# ============================================================
# Read model output
# ============================================================
fin = xr.open_mfdataset(
    FDIR + "out.*.nc",
    combine="by_coords"
)

time = fin.time.sel(time=slice(TSTR, TEND))
X = fin.X.sel(time=slice(TSTR, TEND))
Y = fin.Y.sel(time=slice(TSTR, TEND))
Z = fin.Z.sel(time=slice(TSTR, TEND))

frame_index = list(range(0, len(time), FRAME_SKIP))

if frame_index[-1] != len(time) - 1:
    frame_index.append(len(time) - 1)


# ============================================================
# Animation setup
# ============================================================
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection="3d")

line, = ax.plot([], [], [], color="C0", linewidth=1.0)
point, = ax.plot([], [], [], "o", color="C1", markersize=5)

time_text = ax.text2D(
    0.05,
    0.92,
    "",
    transform=ax.transAxes,
    bbox=dict(facecolor="white", edgecolor="none", alpha=0.8)
)

ax.set_title("Lorenz Climate Model")
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")

ax.set_xlim(-25.0, 25.0)
ax.set_ylim(-30.0, 30.0)
ax.set_zlim(0.0, 60.0)
ax.grid(True)


def init():
    line.set_data([], [])
    line.set_3d_properties([])

    point.set_data([], [])
    point.set_3d_properties([])

    time_text.set_text("")

    return line, point, time_text

def update(frame):

    # -----------------------------------------
    # Normal animation
    # -----------------------------------------
    if frame < len(frame_index):

        iend = frame_index[frame]

        line.set_data(X[:iend+1], Y[:iend+1])
        line.set_3d_properties(Z[:iend+1])

        point.set_data([X[iend]], [Y[iend]])
        point.set_3d_properties([Z[iend]])

        time_text.set_text(f"t = {float(time[iend]):.2f}")

        azim = 30 + 0.20 * frame

    # -----------------------------------------
    # Hold final frame
    # -----------------------------------------
    else:

        iend = frame_index[-1]

        line.set_data(X, Y)
        line.set_3d_properties(Z)

        point.set_data([X[-1]], [Y[-1]])
        point.set_3d_properties([Z[-1]])

        time_text.set_text(f"t = {float(time[-1]):.2f}")

        # Continue rotating slowly
        azim = (
            30
            + 0.20 * len(frame_index)
            + 0.5 * (frame - len(frame_index))
        )

    ax.view_init(elev=25, azim=azim)

    return line, point, time_text

ani = animation.FuncAnimation(
    fig,
    update,
    frames=len(frame_index) + HOLD_FRAMES,
    init_func=init,
    interval=80,
    blit=False
)



# ============================================================
# Save animation and final figure
# ============================================================
plt.tight_layout()

writer = animation.FFMpegWriter(fps=10, bitrate=800)

nframes = len(frame_index) + HOLD_FRAMES

ani.save(
    outfile_mp4,
    writer=writer,
    progress_callback=lambda i, n: print(
        f"\rWriting frame {i+1:4d}/{nframes}",
        end=""
    )
)

print("\nAnimation saved as:", outfile_mp4)

print("Animation saved as:", outfile_mp4)

update(len(frame_index) - 1)
plt.savefig(outfile_png, dpi=150)
print("Final figure saved as:", outfile_png)

plt.show()


#outfile_gif = script_name + ".gif"
#os.system( f'ffmpeg -i {outfile_mp4} -vf "fps=10,scale=900:-1" {outfile_gif}')

fin.close()