{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Fake Data\n",
    "\n",
    "This notebook demonstrates how to create fake data using the functions in `fake_data/fake_data_creator.py`. The data set matches the one created by the `make_demo_data()` in `fake_data/demo_helper.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from kbmod.fake_data.fake_data_creator import *\n",
    "from kbmod.search import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a fake stack of images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the characteristics of the fake data.\n",
    "img_width = 256\n",
    "img_height = 256\n",
    "num_times = 20\n",
    "\n",
    "# Create the fake images\n",
    "fake_times = create_fake_times(num_times, t0=57130.2)\n",
    "ds = FakeDataSet(img_width, img_height, fake_times)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Insert fake moving objects\n",
    "\n",
    "This function creates a random moving object with a given flux that stays within the image for the entire time.  The trajectory is defined by starting pixels (x, y) and velocities (x_v, y_v) of pixels per day."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trj = ds.insert_random_object(500)\n",
    "print(f\"x={trj.x}, y={trj.y}, xv={trj.vx}, yv={trj.vy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can print the object's location at each time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the object was inserted correctly.\n",
    "t0 = ds.stack.get_single_image(0).get_obstime()\n",
    "for i in range(ds.stack.img_count()):\n",
    "    ti = ds.stack.get_single_image(i).get_obstime()\n",
    "    dt = ti - t0\n",
    "    px = int(trj.x + dt * trj.vx + 0.5)\n",
    "    py = int(trj.y + dt * trj.vy + 0.5)\n",
    "\n",
    "    print(f\"{i}: t={ti:.3f} at ({px}, {py})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also create a known trajectory and insert an object using that trajectory. Here we create an object starting at x=50, y=10 and moving at 10 pixels per day in x and 0 pixels per day in y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trj2 = Trajectory(x=50, y=40, vx=10, vy=0, flux=500)\n",
    "ds.insert_object(trj2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot the first and last images and we clearly see our two objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(ds.stack.get_single_image(0).get_science().image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(ds.stack.get_single_image(num_times - 1).get_science().image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save the fake image files\n",
    "\n",
    "We can store the images in a `WorkUnit` and save that to disk (to be used in other notebooks).  In the `WorkUnit` we store a configuration with the minimal search parameters needed to retrieve the object we inserted. Specifically, the data has a known object moving at x_v=10 px/day and y_v = 0 px/day, so we set the search search velocities to [0, 20) and the search angles to [-0.5, 0.5). Note that these bounds may not find the other (random) trajectory we inserted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from kbmod.configuration import SearchConfiguration\n",
    "\n",
    "settings = {\n",
    "    # Override the search data to match the known object.\n",
    "    \"generator_config\": {\n",
    "        \"name\": \"EclipticCenteredSearch\",\n",
    "        \"angles\": [-0.5, 0.5, 11],\n",
    "        \"velocities\": [0.0, 20.0, 21],\n",
    "        \"angle_units\": \"radian\",\n",
    "        \"given_ecliptic\": 0.0,\n",
    "    },\n",
    "    # Loosen the other filtering parameters.\n",
    "    \"clip_negative\": True,\n",
    "    \"sigmaG_lims\": [15, 60],\n",
    "}\n",
    "config = SearchConfiguration.from_dict(settings)\n",
    "\n",
    "# We comment out this line so as not to overwrite the existing demo data.\n",
    "# ds.save_fake_data_to_work_unit(\"../data/demo_data.fits\", config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Jeremy's KBMOD",
   "language": "python",
   "name": "kbmod_jk"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}