{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba83c35d",
   "metadata": {},
   "source": [
    "# Scrambling example\n",
    "\n",
    "In this example we demonstrate how to use the data scrambling mechanism of SkyLLH."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "52af3b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports and helper functions for examples.\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from skyllh.core.config import Config\n",
    "from skyllh.core.dataset import Dataset\n",
    "from skyllh.core.livetime import Livetime\n",
    "from skyllh.core.random import RandomStateService\n",
    "from skyllh.core.scrambling import DataScrambler, UniformRAScramblingMethod\n",
    "from skyllh.core.times import LivetimeTimeGenerationMethod, TimeGenerator\n",
    "from skyllh.i3.scrambling import I3TimeScramblingMethod\n",
    "\n",
    "# Create default configuration.\n",
    "cfg = Config()\n",
    "dataset = Dataset(\n",
    "    cfg=cfg,\n",
    "    name='example',\n",
    "    exp_pathfilenames=None,\n",
    "    mc_pathfilenames=None,\n",
    "    livetime=None,\n",
    "    default_sub_path_fmt='',\n",
    "    version=1,\n",
    ")\n",
    "\n",
    "\n",
    "def gen_data(rss, N):\n",
    "    \"\"\"Create uniformly distributed data on sphere.\"\"\"\n",
    "    arr = np.empty(\n",
    "        (N,),\n",
    "        dtype=[\n",
    "            ('azi', np.float64),\n",
    "            ('zen', np.float64),\n",
    "            ('ra', np.float64),\n",
    "            ('dec', np.float64),\n",
    "            ('time', np.float64),\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    arr['ra'] = rss.random.uniform(0.0, 2.0 * np.pi, N)\n",
    "    arr['dec'] = rss.random.uniform(-np.pi, np.pi, N)\n",
    "\n",
    "    return arr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "822fef84",
   "metadata": {},
   "source": [
    "## Example 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6d3cef5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example 1\n",
      "=========\n",
      "before scrambling: data[\"ra\"]=[2.62022653e+00 4.52593227e+00 7.18638172e-04 1.89961158e+00\n",
      " 9.22094457e-01 5.80180502e-01 1.17030742e+00 2.17122208e+00\n",
      " 2.49296356e+00 3.38548539e+00]\n",
      "after scrambling: data[\"ra\"]=[5.03122651 6.08376691 1.96930219 4.34999129 5.50651545 5.62097944\n",
      " 0.53434854 0.24538844 1.067076   5.51753208]\n"
     ]
    }
   ],
   "source": [
    "def ex1():\n",
    "    \"\"\"Data scrambling via right-ascention scrambling.\"\"\"\n",
    "    print('Example 1')\n",
    "    print('=========')\n",
    "\n",
    "    rss = RandomStateService(seed=1)\n",
    "\n",
    "    # Generate some pseudo data.\n",
    "    data = gen_data(rss=rss, N=10)\n",
    "    print(f'before scrambling: data[\"ra\"]={data[\"ra\"]}')\n",
    "\n",
    "    # Create DataScrambler instance with uniform RA scrambling.\n",
    "    scrambler = DataScrambler(method=UniformRAScramblingMethod())\n",
    "\n",
    "    # Scramble the data.\n",
    "    scrambler.scramble_data(rss=rss, dataset=dataset, data=data)\n",
    "\n",
    "    print(f'after scrambling: data[\"ra\"]={data[\"ra\"]}')\n",
    "\n",
    "\n",
    "ex1()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0df6aaa6",
   "metadata": {},
   "source": [
    "## Example 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f3005ce0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example 2\n",
      "=========\n",
      "before scrambling: data[\"ra\"]=[2.62022653e+00 4.52593227e+00 7.18638172e-04 1.89961158e+00\n",
      " 9.22094457e-01 5.80180502e-01 1.17030742e+00 2.17122208e+00\n",
      " 2.49296356e+00 3.38548539e+00]\n",
      "after scrambling: data[\"ra\"]=[1.95150284 0.42176583 4.80220088 0.7701056  5.19938553 3.15038276\n",
      " 4.77616435 3.81213287 5.62496285 2.56997617]\n"
     ]
    }
   ],
   "source": [
    "def ex2():\n",
    "    \"\"\"Data scrambling via detector on-time scrambling.\"\"\"\n",
    "    print('Example 2')\n",
    "    print('=========')\n",
    "\n",
    "    rss = RandomStateService(seed=1)\n",
    "\n",
    "    # Generate some pseudo data.\n",
    "    data = gen_data(rss=rss, N=10)\n",
    "    print(f'before scrambling: data[\"ra\"]={data[\"ra\"]}')\n",
    "\n",
    "    # Create a Livetime object, which defines the detector live-time.\n",
    "    lt = Livetime(uptime_mjd_intervals_arr=np.array([[55000, 56000], [60000, 69000]], dtype=np.float64))\n",
    "\n",
    "    # Create a TimeGenerator with an on-time time generation method.\n",
    "    timegen = TimeGenerator(method=LivetimeTimeGenerationMethod(livetime=lt))\n",
    "\n",
    "    # Create DataScrambler with IceCube time scrambing method.\n",
    "    scrambler = DataScrambler(method=I3TimeScramblingMethod(timegen))\n",
    "\n",
    "    # Scramble the data.\n",
    "    scrambler.scramble_data(rss=rss, dataset=dataset, data=data)\n",
    "\n",
    "    print(f'after scrambling: data[\"ra\"]={data[\"ra\"]}')\n",
    "\n",
    "\n",
    "ex2()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "skyllh",
   "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.14.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}