{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FVCOM restart from native NEMO AMM7 output\n\n",
    "This tutorial creates a restart file from native NEMO AMM7 physical and ERSEM output. It uses the shared sample data and a short test period.\n\n",
    "The restart is based on a donor FVCOM restart file, which supplies the dimensions and NetCDF variable definitions. The supplied donor is physical-only, so ERSEM writing is disabled by default. Use an ERSEM-capable donor to write tracers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Imports and paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n\n",
    "from pyfvcom2.fvcom_reader import FVCOMReader\n",
    "from pyfvcom2.interpolation import NEMOInterpolator\n",
    "from pyfvcom2.nemo_reader import NEMOReader, default_fvcom_to_nemo_var_names\n",
    "from pyfvcom2.restart import write_restart\n",
    "from pyfvcom2.exceptions import PyFVCOM2ValueError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = Path(os.path.expanduser('~/data/pyfvcom2_doc'))\n",
    "nemo_dir = data_dir / 'AMM7_NEMO'\n",
    "fvcom_dir = data_dir / 'FVCOM_tamar_estuary'\n",
    "donor_filepath = fvcom_dir / 'tamar_v2_donor_restart.nc'\n",
    "output_dir = data_dir / 'restart_from_nemo'\n",
    "output_dir.mkdir(parents=True, exist_ok=True)\n\n",
    "nemo_files = {\n",
    "    'T': nemo_dir / 'amm7_1d_19820101_19820131_grid_T.nc',\n",
    "    'U': nemo_dir / 'amm7_1d_19820101_19820131_grid_U.nc',\n",
    "    'V': nemo_dir / 'amm7_1d_19820101_19820131_grid_V.nc',\n",
    "    'W': nemo_dir / 'amm7_1d_19820101_19820131_grid_W.nc',\n",
    "    'PTRC': nemo_dir / 'amm7_1d_19820101_19820131_ptrc_T.nc',\n",
    "}\n\n",
    "for path in [donor_filepath, *nemo_files.values()]:\n",
    "    print(('OK' if path.exists() else 'MISSING'), path)\n",
    "    if not path.exists():\n",
    "        raise FileNotFoundError(path)\n",
    "if shutil.which('ncks') is None:\n",
    "    raise RuntimeError('write_restart requires the NCO command ncks')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Read the donor and NEMO data\n\n",
    "The PTRC logical grid contains the ERSEM tracers and uses the same coordinate names as the NEMO T grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_date_time = datetime(1982, 1, 1, 12)\n",
    "include_biogeochemistry = False\n\n",
    "fvcom_reader = FVCOMReader(str(donor_filepath))\n",
    "nemo_reader = NEMOReader(\n",
    "    nemo_files,\n",
    "    grid_var_names={'PTRC': {\n",
    "        'time': 'time_counter', 'lon': 'nav_lon',\n",
    "        'lat': 'nav_lat', 'depth': 'deptht',\n",
    "    }},\n",
    ")\n",
    "print('NEMO grids:', nemo_reader.grid_names)\n",
    "for grid in ('T', 'U', 'V', 'PTRC'):\n",
    "    print(grid, nemo_reader.time_span(grid))\n",
    "    if not nemo_reader.contains_date(start_date_time, grid):\n",
    "        raise ValueError(f'{start_date_time} is outside the {grid} range')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Select variables and validate the donor\n\n",
    "The selected ERSEM variables are examples from ptrc_T. Benthic variables require a separate benthos source and matching variables in the donor restart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fvcom_to_nemo = dict(default_fvcom_to_nemo_var_names)\n",
    "physical_variables = ['zeta', 'u', 'v', 'temp', 'salinity']\n",
    "ersem_3d_variables = ['N1_p', 'N3_n', 'O2_o', 'P1_Chl']\n",
    "benthic_variables = []\n",
    "donor_variables = set(fvcom_reader._metadata_dataset.variables)\n",
    "missing_physical = [v for v in physical_variables if v not in donor_variables]\n",
    "if missing_physical:\n",
    "    raise RuntimeError(f'Donor is missing physical variables: {missing_physical}')\n",
    "if include_biogeochemistry:\n",
    "    requested = ersem_3d_variables + benthic_variables\n",
    "    missing = [v for v in requested if v not in donor_variables]\n",
    "    if missing:\n",
    "        raise RuntimeError('Use an ERSEM-capable donor; missing: ' + str(missing))\n",
    "    for variable in requested:\n",
    "        fvcom_to_nemo[variable] = variable\n",
    "print('Donor contains physical restart variables')\n",
    "\n",
    "def source_variable_available(mapping):\n",
    "    candidates = (mapping,) if isinstance(mapping, str) else mapping\n",
    "    return any(\n",
    "        name in dataset.variables\n",
    "        for name in candidates\n",
    "        for dataset in nemo_reader._metadata_datasets.values()\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Interpolate physical and optional ERSEM fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpolator = NEMOInterpolator(nemo_reader, fvcom_to_nemo_var_names=fvcom_to_nemo)\n",
    "node_coords = fvcom_reader.get_interpolation_coordinates('node', 'layer_centre')\n",
    "node_coords.dates = np.asarray([start_date_time])\n",
    "element_coords = fvcom_reader.get_interpolation_coordinates('element', 'layer_centre')\n",
    "element_coords.dates = np.asarray([start_date_time])\n",
    "positions = {\n",
    "    'temp': node_coords, 'salinity': node_coords,\n",
    "    'u': element_coords, 'v': element_coords,\n",
    "}\n",
    "if source_variable_available(fvcom_to_nemo['zeta']):\n",
    "    positions['zeta'] = node_coords\n",
    "else:\n",
    "    print('NEMO zeta is unavailable; donor zeta will be retained')\n",
    "if include_biogeochemistry:\n",
    "    positions.update({v: node_coords for v in ersem_3d_variables})\n",
    "    positions.update({v: node_coords for v in benthic_variables})\n",
    "var_data = {}\n",
    "for variable, coordinates in positions.items():\n",
    "    try:\n",
    "        values = interpolator.interpolate(coordinates, variable)\n",
    "    except Exception as exc:\n",
    "        raise RuntimeError(f'Interpolation failed for {variable}: {exc}') from exc\n",
    "    var_data[variable] = values[0] if values.ndim > 0 and values.shape[0] == 1 else values\n",
    "    print(variable, values.shape, '->', var_data[variable].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Write the restart file\n\n",
    "`write_restart` preserves the donor structure and replaces the selected variables. It requires the NCO `ncks` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "restart_file = output_dir / f'nsea_nemo_restart_{start_date_time:%Y%m%d}.nc'\n",
    "write_restart(str(donor_filepath), str(restart_file), var_data, start_date_time)\n",
    "print('Created:', restart_file)\n",
    "print('Size (MB):', restart_file.stat().st_size / 1024**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inspect and visualise the restart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netCDF4 import Dataset\n",
    "import matplotlib.pyplot as plt\n\n",
    "with Dataset(restart_file) as ds:\n",
    "    temp = np.asarray(ds.variables['temp'][0])\n",
    "    u = np.asarray(ds.variables['u'][0])\n",
    "    v = np.asarray(ds.variables['v'][0])\n",
    "    temp_surface = temp[-1, :] if temp.ndim == 2 else temp\n",
    "    u_surface = u[-1, :] if u.ndim == 2 else u\n",
    "    v_surface = v[-1, :] if v.ndim == 2 else v\n",
    "    speed = np.hypot(u_surface, v_surface)\n\n",
    "fig, axes = plt.subplots(1, 2, figsize=(13, 5), constrained_layout=True)\n",
    "a = axes[0].scatter(fvcom_reader.lon_nodes, fvcom_reader.lat_nodes, c=temp_surface, s=5, cmap='plasma')\n",
    "axes[0].set_title('Restart surface temperature')\n",
    "axes[0].set_xlabel('Longitude'); axes[0].set_ylabel('Latitude')\n",
    "fig.colorbar(a, ax=axes[0], label='degC')\n",
    "b = axes[1].scatter(fvcom_reader.lon_elements, fvcom_reader.lat_elements, c=speed, s=5, cmap='magma')\n",
    "axes[1].set_title('Restart surface current speed')\n",
    "axes[1].set_xlabel('Longitude'); axes[1].set_ylabel('Latitude')\n",
    "fig.colorbar(b, ax=axes[1], label='m/s')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extending to a full FVCOM-ERSEM restart\n\n",
    "Use an ERSEM-capable donor restart, set `include_biogeochemistry = True`, and populate `benthic_variables` with fields available in both the NEMO benthos output and donor. The notebook deliberately stops when a requested variable is missing instead of silently retaining stale donor values."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}