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


# ============================================================
# Figure settings
# ============================================================
script_name = os.path.splitext(os.path.basename(__file__))[0]
outfile_png = script_name + ".png"


# ============================================================
# Input files
# ============================================================
FDIR = "./out/"
OBSFILE = "observation.nc"

TSTR = 50
TEND = 70

# ============================================================
# Read synthetic observations
# ============================================================
fobs = xr.open_dataset(OBSFILE)

otime = fobs.time

Xtrue = fobs.Xtrue
Ytrue = fobs.Ytrue
Ztrue = fobs.Ztrue

Xobs = fobs.Xobs
Yobs = fobs.Yobs
Zobs = fobs.Zobs


# ============================================================
# Read model output
# ============================================================
#fin = xr.open_mfdataset(
#    FDIR + "out.*.nc",
#    combine="by_coords"
#)
#
#time = fin.time.sel(time=slice(TSTR, TEND))
#
#Xmdl = fin.X.sel(time=slice(TSTR, TEND))
#Ymdl = fin.Y.sel(time=slice(TSTR, TEND))
#Zmdl = fin.Z.sel(time=slice(TSTR, TEND))


# ============================================================
# Plot
# ============================================================
fig, ax = plt.subplots(3, 1, figsize=(9, 7), sharex=True)

# ------------------------------------------------------------
# X
# ------------------------------------------------------------
#ax[0].plot( time, Xmdl, color="gray", linewidth=1.0, label="Model trajectory")
ax[0].plot( otime, Xtrue, "-x",  color="C1", markersize=8, markeredgewidth=2.5, label="True state")
ax[0].plot( otime, Xobs, "-o",  markerfacecolor="none", markeredgecolor="C0", markeredgewidth=2.5, markersize=7, label="Observation")

ax[0].set_ylabel("X")
ax[0].set_title("Synthetic Observations")
ax[0].legend(loc="best", fontsize=8)


# ------------------------------------------------------------
# Y
# ------------------------------------------------------------
#ax[1].plot( time, Ymdl, color="gray", linewidth=1.0)
ax[1].plot( otime, Ytrue, "-x",  color="C1", markersize=8, markeredgewidth=2.5)
ax[1].plot( otime, Yobs, "-o",  markerfacecolor="none", markeredgecolor="C0", markeredgewidth=2.5, markersize=7)

ax[1].set_ylabel("Y")


# ------------------------------------------------------------
# Z
# ------------------------------------------------------------
#ax[2].plot( time, Zmdl, color="gray", linewidth=1.0)
ax[2].plot( otime, Ztrue, "-x", color="C1", markersize=8, markeredgewidth=2.5)
ax[2].plot( otime, Zobs, "-o",  markerfacecolor="none", markeredgecolor="C0", markeredgewidth=2.5, markersize=7)

ax[2].set_ylabel("Z")
ax[2].set_xlabel("Time")


# ============================================================
# Figure settings
# ============================================================
ax[0].set_xlim(TSTR, TEND+1)

ax[0].set_ylim(-25.0, 25.0)
ax[1].set_ylim(-25.0, 25.0)
ax[2].set_ylim(0.0, 60.0)

for a in ax:
    a.grid(True)

plt.tight_layout()
plt.savefig(outfile_png, dpi=150)
plt.show()
plt.close()

print("Figure saved as:", outfile_png)


# ============================================================
# Close files
# ============================================================
fobs.close()
#fin.close()