import os
import argparse
import numpy as np
import xarray as xr


# ============================================================
# Default settings
# ============================================================
OBSFILE = "../../CTL/observation.nc"

DEFAULT_T_INITIAL = 55

RDIR = "./rest/"


# ============================================================
# Command-line arguments
# ============================================================
parser = argparse.ArgumentParser(
    description=(
        "Create a model restart file using synthetic observations "
        "as the initial condition."
    )
)

parser.add_argument(
    "--t_initial",
    type=int,
    default=DEFAULT_T_INITIAL,
    help=(
        "Forecast initialization time "
        f"(default: {DEFAULT_T_INITIAL})"
    ),
)

args = parser.parse_args()

T_INITIAL = args.t_initial

RFILE = f"restart.{T_INITIAL:04d}.nc"


# ============================================================
# Create output directory
# ============================================================
os.makedirs(RDIR, exist_ok=True)


# ============================================================
# Read observation file
# ============================================================
if not os.path.exists(OBSFILE):
    raise FileNotFoundError(
        f"Observation file was not found: {OBSFILE}"
    )

fin = xr.open_dataset(OBSFILE)

if T_INITIAL not in fin.time.values:
    available_times = fin.time.values

    fin.close()

    raise ValueError(
        f"Initialization time {T_INITIAL} is not available in {OBSFILE}.\n"
        f"Available observation times are:\n{available_times}"
    )

obs = fin.sel(time=T_INITIAL)


# ============================================================
# Extract observed initial state
# ============================================================
time = np.float64(T_INITIAL)

Xint_value = obs.Xobs.item()
Yint_value = obs.Yobs.item()
Zint_value = obs.Zobs.item()

Xint = xr.DataArray(
    [Xint_value],
    dims=["time"],
    coords={"time": [time]},
)

Yint = xr.DataArray(
    [Yint_value],
    dims=["time"],
    coords={"time": [time]},
)

Zint = xr.DataArray(
    [Zint_value],
    dims=["time"],
    coords={"time": [time]},
)


# ============================================================
# External forcing state
# ============================================================
Xext = xr.DataArray(
    [0.0],
    dims=["time"],
    coords={"time": [time]},
)

Yext = xr.DataArray(
    [0.0],
    dims=["time"],
    coords={"time": [time]},
)

Zext = xr.DataArray(
    [0.0],
    dims=["time"],
    coords={"time": [time]},
)


# ============================================================
# Create restart dataset
# ============================================================
dsout = xr.Dataset(
    {
        "Xint": Xint,
        "Yint": Yint,
        "Zint": Zint,
        "Xext": Xext,
        "Yext": Yext,
        "Zext": Zext,
    },
    attrs={
        "description": (
            "Model restart file initialized from synthetic observations"
        ),
        "observation_file": OBSFILE,
        "initial_time": T_INITIAL,
    },
)


# ============================================================
# Save restart file
# ============================================================
outfile = os.path.join(RDIR, RFILE)

dsout.to_netcdf(outfile)

dsout.close()
fin.close()


# ============================================================
# Print summary
# ============================================================
print("--------------------------------")
print("Create initial condition")
print("--------------------------------")
print("Observation file :", OBSFILE)
print("Initial time     :", T_INITIAL)
print("Restart file     :", outfile)
print("")
print("Initial state")
print(f"X = {Xint_value:10.4f}")
print(f"Y = {Yint_value:10.4f}")
print(f"Z = {Zint_value:10.4f}")