{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import supporting package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from DataContainer.ReadData import read_hdf5_file\n",
    "from Analyser.ImagingAnalyser import ImageAnalyser\n",
    "from Analyser.FitAnalyser import FitAnalyser\n",
    "\n",
    "imageAnalyser = ImageAnalyser()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start a client for parallel computing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "client = Client(n_workers=6, threads_per_worker=10, processes=True, memory_limit='10GB')\n",
    "client"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set file path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filepath = \"//DyLabNAS/Data/Evaporative_Cooling/2023/05/03/0043/*.h5\"\n",
    "# filepath = \"//DyLabNAS/Data/Evaporative_Cooling/2023/04/18/0003/2023-04-18_0003_Evaporative_Cooling_000.h5\"\n",
    "\n",
    "filepath = \"//DyLabNAS/Data/Repetition_scan/2023/04/21/0002/*.h5\"\n",
    "\n",
    "# filepath = r\"./testData/0000/*.h5\"\n",
    "\n",
    "# filepath = \"//DyLabNAS/Data/Evaporative_Cooling/2023/04/18/0003/*.h5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "groupList = [\n",
    "    \"images/MOT_3D_Camera/in_situ_absorption\",\n",
    "    # \"images/ODT_1_Axis_Camera/in_situ_absorption\",\n",
    "]\n",
    "\n",
    "dskey = {\n",
    "    \"images/MOT_3D_Camera/in_situ_absorption\": \"camera_1\",\n",
    "    # \"images/ODT_1_Axis_Camera/in_situ_absorption\": \"camera_2\",\n",
    "}\n",
    "\n",
    "datasetDict = {\n",
    "    dskey[groupList[i]]: read_hdf5_file(filepath, groupList[i])\n",
    "    for i in range(len(groupList))\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rechunk the data for parallel computing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasetDict[\"camera_1\"]\n",
    "\n",
    "scanAxis = dataset.scanAxis\n",
    "\n",
    "# dataset = dataset.chunk(\n",
    "#     {\n",
    "#         # \"compZ_current_sg\": \"auto\",\n",
    "#         \"sin_mod_freq\": \"auto\",\n",
    "#         \"runs\": 2,\n",
    "#         \"x\": \"auto\",\n",
    "#         \"y\": \"auto\",\n",
    "#     }\n",
    "# )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate absorption imaging"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## get OD images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = imageAnalyser.get_absorption_images(dataset)\n",
    "\n",
    "dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select region of interests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageAnalyser.center = (529, 962)\n",
    "imageAnalyser.span = (100,100)\n",
    "imageAnalyser.fraction = (0.1, 0.1)\n",
    "\n",
    "# imageAnalyser.center = (890, 1150)\n",
    "# imageAnalyser.span = (600,600)\n",
    "# imageAnalyser.fraction = (0.1, 0.1)\n",
    "\n",
    "# imageAnalyser.center = (890, 950)\n",
    "# imageAnalyser.span = (100,100)\n",
    "# imageAnalyser.fraction = (0.1, 0.1)\n",
    "\n",
    "dataset_crop = imageAnalyser.crop_image(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_crop.OD.isel(runs=[0]).plot.pcolormesh(cmap='jet', vmin=0, col=scanAxis[0], row=scanAxis[1])\n",
    "# dataset_crop.OD.plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove the background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_crop['OD'] = dataset_crop['OD'] + 500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_crop['OD'] = imageAnalyser.substract_offset(dataset_crop['OD'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitAnalyser = FitAnalyser(\"Two Gaussian-2D\", fitDim=2)\n",
    "\n",
    "params = fitAnalyser.guess(dataset_crop.OD, dask=\"parallelized\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitResult = fitAnalyser.fit(dataset_crop.OD, params, dask=\"parallelized\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitCurve = fitAnalyser.eval(fitResult, x=np.arange(100), y=np.arange(100)).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitCurve.isel(**{scanAxis[0]:np.arange(30), 'runs':range(dataset_crop.OD['runs'].size)}).plot.pcolormesh(cmap='jet', vmin=0, col=scanAxis[0], row=scanAxis[1])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get the Ncount"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ncount = dataset_crop.OD.sum(dim=(scanAxis[1], 'x', 'y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "Ncount.plot(fig=fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c05913ad4f24fdc6b2418069394dc5835b1981849b107c9ba6df693aafd66650"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}