{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Extraction\n",
    "\n",
    "This notebook demonstrates gcmprocpy's data extraction functions (`arr_*`) for slicing model output along various dimensions, as well as height interpolation utilities.\n",
    "\n",
    "**Note:** This notebook requires TIE-GCM or WACCM-X model output files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using time: 2002-12-21T01:00:00.000000000\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import numpy as np\n",
    "import gcmprocpy as gy\n",
    "\n",
    "directory = '/glade/work/nikhilr/tiegcm3.0/benchmarks/2.5/seasons/decsol_smin/hist'\n",
    "datasets = gy.load_datasets(directory, dataset_filter='sech')\n",
    "times = gy.time_list(datasets)\n",
    "t_val = times[0]\n",
    "print(f'Using time: {t_val}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lat_lon — Latitude x Longitude Slice\n",
    "\n",
    "Extract a 2D lat/lon array at a given time and pressure level. Use `plot_mode=True` to get a `PlotData` object with metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape:  (72, 144)\n",
      "unit:   K\n",
      "level:  5.0\n",
      "lats:   -88.8 to 88.8\n",
      "lons:   -180.0 to 177.5\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lat_lon(datasets, 'TN', t_val, selected_lev_ilev=5.0, plot_mode=True)\n",
    "print(f'shape:  {result.values.shape}')\n",
    "print(f'unit:   {result.variable_unit}')\n",
    "print(f'level:  {result.selected_lev}')\n",
    "print(f'lats:   {result.lats[0]:.1f} to {result.lats[-1]:.1f}')\n",
    "print(f'lons:   {result.lons[0]:.1f} to {result.lons[-1]:.1f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_var — Full Variable at a Time\n",
    "\n",
    "Extract the full 3D field (lev x lat x lon) at a single time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (57, 72, 144)  (nlev x nlat x nlon)\n",
      "levs:  (57,)\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_var(datasets, 'TN', t_val, plot_mode=True)\n",
    "print(f'shape: {result.values.shape}  (nlev x nlat x nlon)')\n",
    "print(f'levs:  {result.levs.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lev_var — Vertical Profile\n",
    "\n",
    "Extract a 1D vertical profile at a specific latitude and longitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (57,)\n",
      "levs:  [-6.875 -6.625 -6.375 -6.125 -5.875 -5.625 -5.375 -5.125 -4.875 -4.625\n",
      " -4.375 -4.125 -3.875 -3.625 -3.375 -3.125 -2.875 -2.625 -2.375 -2.125\n",
      " -1.875 -1.625 -1.375 -1.125 -0.875 -0.625 -0.375 -0.125  0.125  0.375\n",
      "  0.625  0.875  1.125  1.375  1.625  1.875  2.125  2.375  2.625  2.875\n",
      "  3.125  3.375  3.625  3.875  4.125  4.375  4.625  4.875  5.125  5.375\n",
      "  5.625  5.875  6.125  6.375  6.625  6.875  7.125]\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lev_var(datasets, 'TN', t_val, selected_lat=0.0, selected_lon=0.0, plot_mode=True)\n",
    "print(f'shape: {result.values.shape}')\n",
    "print(f'levs:  {result.levs}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lev_lon — Level x Longitude Cross-Section\n",
    "\n",
    "Extract a 2D lev/lon cross-section at a given latitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (57, 144)  (nlev x nlon)\n",
      "levs:  (57,)\n",
      "lons:  (144,)\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lev_lon(datasets, 'TN', t_val, selected_lat=0.0, plot_mode=True)\n",
    "print(f'shape: {result.values.shape}  (nlev x nlon)')\n",
    "print(f'levs:  {result.levs.shape}')\n",
    "print(f'lons:  {result.lons.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lev_lat — Level x Latitude Cross-Section\n",
    "\n",
    "Extract a 2D lev/lat cross-section at a given longitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (57, 72)  (nlev x nlat)\n",
      "levs:  (57,)\n",
      "lats:  (72,)\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lev_lat(datasets, 'TN', t_val, selected_lon=0.0, plot_mode=True)\n",
    "print(f'shape: {result.values.shape}  (nlev x nlat)')\n",
    "print(f'levs:  {result.levs.shape}')\n",
    "print(f'lats:  {result.lats.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lev_time — Level x Time Cross-Section\n",
    "\n",
    "Extract a 2D lev/time cross-section at a specific latitude and longitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape:  (57, 120)  (nlev x ntime)\n",
      "levs:   (57,)\n",
      "mtimes: 120\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lev_time(datasets, 'TN', selected_lat=0.0, selected_lon=0.0, plot_mode=True)\n",
    "print(f'shape:  {result.values.shape}  (nlev x ntime)')\n",
    "print(f'levs:   {result.levs.shape}')\n",
    "print(f'mtimes: {len(result.mtime_values)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lat_time — Latitude x Time Cross-Section\n",
    "\n",
    "Extract a 2D lat/time cross-section at a given longitude and level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape:  (72, 120)  (nlat x ntime)\n",
      "lats:   (72,)\n",
      "mtimes: 120\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lat_time(datasets, 'TN', selected_lon=0.0, selected_lev_ilev=5.0, plot_mode=True)\n",
    "print(f'shape:  {result.values.shape}  (nlat x ntime)')\n",
    "print(f'lats:   {result.lats.shape}')\n",
    "print(f'mtimes: {len(result.mtime_values)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_lon_time — Longitude x Time Cross-Section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape:  (144, 120)  (nlon x ntime)\n",
      "lons:   (144,)\n",
      "mtimes: 120\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_lon_time(datasets, 'TN', selected_lat=0.0, selected_lev_ilev=5.0, plot_mode=True)\n",
    "print(f'shape:  {result.values.shape}  (nlon x ntime)')\n",
    "print(f'lons:   {result.lons.shape}')\n",
    "print(f'mtimes: {len(result.mtime_values)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_var_time — Variable vs Time Series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape:  (120,)  (ntime,)\n",
      "mtimes: 120\n"
     ]
    }
   ],
   "source": [
    "result = gy.arr_var_time(datasets, 'TN', selected_lat=0.0, selected_lon=0.0, selected_lev_ilev=5.0, plot_mode=True)\n",
    "print(f'shape:  {result.values.shape}  (ntime,)')\n",
    "print(f'mtimes: {len(result.mtime_values)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## batch_arr_lat_lon — Multiple Variables at Once\n",
    "\n",
    "Extract several variables in a single pass over the datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TN: shape=(72, 144), unit=K\n",
      "O1: shape=(72, 144), unit=mmr\n",
      "NO: shape=(72, 144), unit=mmr\n"
     ]
    }
   ],
   "source": [
    "results = gy.batch_arr_lat_lon(datasets, ['TN', 'O1', 'NO'], t_val, selected_lev_ilev=5.0, plot_mode=True)\n",
    "for name, r in results.items():\n",
    "    print(f'{name}: shape={r.values.shape}, unit={r.variable_unit}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## arr_sat_track — Satellite Track Interpolation\n",
    "\n",
    "Interpolate model data along a satellite trajectory using trilinear interpolation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "At fixed level: shape=(20,)\n",
      "All levels:     shape=(57, 20)\n"
     ]
    }
   ],
   "source": [
    "# Simulate a LEO ascending pass\n",
    "sat_time = np.array([times[0] + np.timedelta64(i * 6, 'm') for i in range(20)])\n",
    "sat_lat = np.linspace(-60, 60, 20)\n",
    "sat_lon = np.linspace(-120, 120, 20)\n",
    "\n",
    "# 1D extraction at a fixed level\n",
    "values = gy.arr_sat_track(datasets, 'NO', sat_time, sat_lat, sat_lon, selected_lev_ilev=5.0)\n",
    "print(f'At fixed level: shape={values.shape}')\n",
    "\n",
    "# 2D extraction across all levels\n",
    "values_2d = gy.arr_sat_track(datasets, 'NO', sat_time, sat_lat, sat_lon)\n",
    "print(f'All levels:     shape={values_2d.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Height Interpolation\n",
    "\n",
    "### height_to_pres_level\n",
    "\n",
    "Convert a target height in km to the nearest pressure level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300 km -> pressure level 3.0\n",
      "300 km at equator, 0E -> pressure level 3.25\n",
      "  100 km -> -6.25\n",
      "  200 km -> 0.0\n",
      "  300 km -> 3.0\n",
      "  400 km -> 5.25\n",
      "  500 km -> 7.0\n"
     ]
    }
   ],
   "source": [
    "pres_level = gy.height_to_pres_level(datasets, t_val, 300.0)\n",
    "print(f'300 km -> pressure level {pres_level}')\n",
    "\n",
    "# At a specific location\n",
    "pres_level_ll = gy.height_to_pres_level(datasets, t_val, 300.0, latitude=0.0, longitude=0.0)\n",
    "print(f'300 km at equator, 0E -> pressure level {pres_level_ll}')\n",
    "\n",
    "# Different heights map to different levels\n",
    "for h in [100.0, 200.0, 300.0, 400.0, 500.0]:\n",
    "    pl = gy.height_to_pres_level(datasets, t_val, h)\n",
    "    print(f'  {h:.0f} km -> {pl}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### interpolate_to_height\n",
    "\n",
    "Interpolate a 2D field from pressure levels to constant height surfaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input shape:  (57, 72) (nlev x nlat)\n",
      "Output shape: (30, 72) (n_heights x nlat)\n",
      "Height range: 96.4 to 482.7 km\n",
      "\n",
      "Custom heights shape: (5, 72)\n"
     ]
    }
   ],
   "source": [
    "# Extract a lev x lat cross-section on pressure levels\n",
    "result_lev_lat = gy.arr_lev_lat(datasets, 'TN', t_val, selected_lon=0.0, plot_mode=True)\n",
    "\n",
    "# Interpolate to 30 evenly spaced height levels\n",
    "interp_vals, target_heights = gy.interpolate_to_height(\n",
    "    datasets, result_lev_lat.values, result_lev_lat.levs, t_val, n_heights=30)\n",
    "print(f'Input shape:  {result_lev_lat.values.shape} (nlev x nlat)')\n",
    "print(f'Output shape: {interp_vals.shape} (n_heights x nlat)')\n",
    "print(f'Height range: {target_heights[0]:.1f} to {target_heights[-1]:.1f} km')\n",
    "\n",
    "# Interpolate to specific heights\n",
    "custom_hts = np.array([100.0, 200.0, 300.0, 400.0, 500.0])\n",
    "interp_custom, _ = gy.interpolate_to_height(\n",
    "    datasets, result_lev_lat.values, result_lev_lat.levs, t_val, target_heights=custom_hts)\n",
    "print(f'\\nCustom heights shape: {interp_custom.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datasets closed.\n"
     ]
    }
   ],
   "source": [
    "gy.close_datasets(datasets)\n",
    "print('Datasets closed.')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gcmprocpy_test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}