import os
import numpy as np
import xarray as xr


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

outfile_nc = "observation.nc"


# ============================================================
# Input model output
# ============================================================
FDIR = "./out/"

TSTR = 50
TEND = 70


# ============================================================
# Observation settings
# ============================================================
OERR = 1.0      # observation-error standard deviation
SEED = 150      # random seed

OBS_TSTR = 55.0
OBS_TEND = 70.0
OBS_DT = 1.0    # observation interval


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

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


# ============================================================
# Select observation times
# ============================================================
otime = np.arange(OBS_TSTR, OBS_TEND + OBS_DT, OBS_DT)


# ============================================================
# Sample true state at observation times
# ============================================================
Xtrue = Xtrue_all.sel(time=otime, method="nearest")
Ytrue = Ytrue_all.sel(time=otime, method="nearest")
Ztrue = Ztrue_all.sel(time=otime, method="nearest")

# Replace coordinates with exact observation times
Xtrue = Xtrue.assign_coords(time=otime)
Ytrue = Ytrue.assign_coords(time=otime)
Ztrue = Ztrue.assign_coords(time=otime)


# ============================================================
# Generate synthetic observations
# ============================================================
np.random.seed(SEED)

Xobs = Xtrue + np.random.normal(0.0, OERR, len(otime))
Yobs = Ytrue + np.random.normal(0.0, OERR, len(otime))
Zobs = Ztrue + np.random.normal(0.0, OERR, len(otime))


# ============================================================
# Save synthetic observations
# ============================================================
dsobs = xr.Dataset(
    {
        "Xtrue": Xtrue,
        "Ytrue": Ytrue,
        "Ztrue": Ztrue,
        "Xobs": Xobs,
        "Yobs": Yobs,
        "Zobs": Zobs,
    },
    coords={
        "time": otime
    },
    attrs={
        "description": "Synthetic observations generated from a perfect-model experiment",
        "observation_error_std": OERR,
        "random_seed": SEED,
        "observation_interval": OBS_DT,
    },
)

dsobs.to_netcdf(outfile_nc)
dsobs.close()

print("Observation file saved as:", outfile_nc)

fin.close()