{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "7520724c-2be0-44a3-b657-4247cd9c497d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import configparser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7b936622-1122-438e-a7d6-c302922aaadc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'#Units: Masses in Msun / h\\n'\n",
      "b'#Units: Positions in Mpc / h (comoving)\\n'\n",
      "b'#Units: Velocities in km / s (physical, peculiar)\\n'\n",
      "b'#Units: Halo Distances, Lengths, and Radii in kpc / h (comoving)\\n'\n",
      "b'#Units: Angular Momenta in (Msun/h) * (Mpc/h) * km/s (physical)\\n'\n",
      "b'#Units: Spins are dimensionless\\n'\n",
      "b'#Units: Total energy in (Msun/h)*(km/s)^2 (physical)\\n'\n"
     ]
    }
   ],
   "source": [
    "dirname = '/data54/lavaux/VELMASS/halo_central/halos_new/'\n",
    "lam = 5\n",
    "R_max = 100\n",
    "Mmin = 3e12\n",
    "    \n",
    "with open(dirname + \"/halos_0.0.ascii\", \"rb\") as f:\n",
    "    for _ in range(30):\n",
    "        r = str(f.readline())\n",
    "        if \"Units\" in r:\n",
    "            print(r)\n",
    "\n",
    "# Get box size\n",
    "with open(dirname + '/auto-rockstar.cfg') as f:\n",
    "    data = [r for r in f]\n",
    "Lbox = [r for r in data if r.startswith('BOX_SIZE')][0].strip()\n",
    "Lbox = float(Lbox.split('=')[1])\n",
    "origin = np.array([Lbox/2, Lbox/2, Lbox/2])    \n",
    "omega_m = [r for r in data if r.startswith('Om')][0].strip()\n",
    "omega_m = float(omega_m.split('=')[1])\n",
    "\n",
    "halo_file = dirname + '/out.npy'\n",
    "halos = np.load(halo_file)\n",
    "\n",
    "# Cut by mass\n",
    "# Mmin = ???\n",
    "halos = halos[halos['Mvir'] > Mmin]\n",
    "\n",
    "xtrue = halos['X']  - origin[0] # Mpc / h\n",
    "ytrue = halos['Y']  - origin[1] # Mpc / h\n",
    "ztrue = halos['Z']  - origin[2] # Mpc / h\n",
    "rtrue = np.sqrt(xtrue**2 + ytrue**2 + ztrue**2) # Mpc / h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "2c5f507e-348f-4e12-904f-3e8dd8455a58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12449803\n",
      "2000.0\n",
      "\tMade 342 of 400\n",
      "\tMade 400 of 400\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x149298041100>]"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31488 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Sample \n",
    "Nt = 400\n",
    "R_lim = 100\n",
    "print(len(rtrue))\n",
    "print(Lbox)\n",
    "\n",
    "# Get relative probability of each object according to selection function\n",
    "# lam = float(config[f'sample_{i}']['lam'])\n",
    "# R_max = float(config['mock']['R_max'])\n",
    "prob = np.exp(- lam * rtrue / R_max)\n",
    "prob /= np.sum(prob)\n",
    "\n",
    "accepted_count = 0\n",
    "rsamp = np.zeros(Nt)\n",
    "\n",
    "# Loop until we have Nt valid positions\n",
    "while accepted_count < Nt:\n",
    "\n",
    "    ids = np.random.choice(len(rtrue), size=Nt, p=prob, replace=False)\n",
    "    \n",
    "    # Here I would actually scatter\n",
    "    \n",
    "    valid_indices = rtrue[ids] < R_lim\n",
    "    ids = ids[valid_indices]\n",
    "    \n",
    "    # Calculate how many valid positions we need to reach Nt\n",
    "    remaining_needed = Nt - accepted_count\n",
    "    selected_count = min(len(ids), remaining_needed)\n",
    "    ids = ids[:selected_count]\n",
    "\n",
    "    # Append only the needed number of valid positions\n",
    "    rsamp[accepted_count:accepted_count+selected_count] = rtrue[ids]\n",
    "\n",
    "    # Update the accepted count\n",
    "    accepted_count += selected_count\n",
    "\n",
    "    print(f'\\tMade {accepted_count} of {Nt}')\n",
    "    \n",
    "    # Set up for next iteration\n",
    "    prob[ids] = 0\n",
    "    prob /= np.sum(prob)\n",
    "\n",
    "plt.hist(rsamp, bins=30, density=True, alpha=0.5)\n",
    "x = plt.gca().get_xlim()\n",
    "x = [max(0, x[0]), x[1]]\n",
    "x = np.linspace(x[0], x[1], 100)\n",
    "y = x ** 2 * np.exp( - lam * x / R_max) * (lam / R_max) ** 3 / 2\n",
    "plt.plot(x, y, color='k')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "73491c07-4ad8-4aa7-b277-192bc28f26d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 12449803)\n"
     ]
    }
   ],
   "source": [
    "comb_x = np.array([xtrue, ytrue, ztrue])\n",
    "print(comb_x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "7302fe50-8271-4d0f-bcf2-dc02bdae4f94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['FILE_FORMAT', 'PARTICLE_MASS', 'MASS_DEFINITION', 'MASS_DEFINITION2', 'MASS_DEFINITION3', 'MASS_DEFINITION4', 'MASS_DEFINITION5', 'STRICT_SO_MASSES', 'MIN_HALO_OUTPUT_SIZE', 'FORCE_RES', 'FORCE_RES_PHYS_MAX', 'SCALE_NOW', 'h0', 'Ol', 'Om', 'W0', 'WA', 'GADGET_ID_BYTES', 'GADGET_MASS_CONVERSION', 'GADGET_LENGTH_CONVERSION', 'GADGET_SKIP_NON_HALO_PARTICLES', 'GADGET_HALO_PARTICLE_TYPE', 'RESCALE_PARTICLE_MASS', 'TIPSY_LENGTH_CONVERSION', 'TIPSY_VELOCITY_CONVERSION', 'PARALLEL_IO', 'PARALLEL_IO_SERVER_ADDRESS', 'PARALLEL_IO_SERVER_PORT', 'PARALLEL_IO_WRITER_PORT', 'PARALLEL_IO_SERVER_INTERFACE', 'PARALLEL_IO_CATALOGS', 'RUN_ON_SUCCESS', 'RUN_PARALLEL_ON_SUCCESS', 'LOAD_BALANCE_SCRIPT', 'INBASE', 'FILENAME', 'STARTING_SNAP', 'RESTART_SNAP', 'NUM_SNAPS', 'NUM_BLOCKS', 'NUM_READERS', 'PRELOAD_PARTICLES', 'SNAPSHOT_NAMES', 'LIGHTCONE_ALT_SNAPS', 'BLOCK_NAMES', 'OUTBASE', 'OVERLAP_LENGTH', 'NUM_WRITERS', 'FORK_READERS_FROM_WRITERS', 'FORK_PROCESSORS_PER_MACHINE', 'OUTPUT_FORMAT', 'DELETE_BINARY_OUTPUT_AFTER_FINISHED', 'FULL_PARTICLE_CHUNKS', 'BGC2_SNAPNAMES', 'SHAPE_ITERATIONS', 'WEIGHTED_SHAPES', 'BOUND_PROPS', 'BOUND_OUT_TO_HALO_EDGE', 'DO_MERGER_TREE_ONLY', 'IGNORE_PARTICLE_IDS', 'EXACT_LL_CALC', 'TRIM_OVERLAP', 'ROUND_AFTER_TRIM', 'LIGHTCONE', 'PERIODIC', 'LIGHTCONE_ORIGIN', 'LIGHTCONE_ALT_ORIGIN', 'LIMIT_CENTER', 'LIMIT_RADIUS', 'SWAP_ENDIANNESS', 'GADGET_VARIANT', 'ART_VARIANT', 'FOF_FRACTION', 'FOF_LINKING_LENGTH', 'INITIAL_METRIC_SCALING', 'INCLUDE_HOST_POTENTIAL_RATIO', 'TEMPORAL_HALO_FINDING', 'MIN_HALO_PARTICLES', 'UNBOUND_THRESHOLD', 'ALT_NFW_METRIC', 'EXTRA_PROFILING', 'TOTAL_PARTICLES', 'BOX_SIZE', 'OUTPUT_LEVELS', 'DUMP_PARTICLES', 'ROCKSTAR_CONFIG_FILENAME', 'AVG_PARTICLE_SPACING', 'SINGLE_SNAP', 'RAMSES_DP'])\n",
      "0.315\n",
      "<class 'float'>\n"
     ]
    }
   ],
   "source": [
    "dirname = '/data54/lavaux/VELMASS/halo_central/halos_new/'\n",
    "\n",
    "def parse_file_to_dict(file_path):\n",
    "    config_dict = {}\n",
    "    with open(file_path, 'r') as file:\n",
    "        for line in file:\n",
    "            line = line.strip()\n",
    "            if not line or line.startswith(\"#\"):  # Skip empty lines and comments\n",
    "                continue\n",
    "            key, value = line.split(\"=\", 1)\n",
    "            key = key.strip()\n",
    "            value = value.strip()\n",
    "\n",
    "            # Convert the value to the appropriate type (int, float, or leave as string)\n",
    "            if value.isdigit():\n",
    "                value = int(value)\n",
    "            else:\n",
    "                try:\n",
    "                    value = float(value)\n",
    "                except ValueError:\n",
    "                    pass  # Leave it as a string if it cannot be converted\n",
    "\n",
    "            config_dict[key] = value\n",
    "\n",
    "    return config_dict\n",
    "\n",
    "\n",
    "config_dict = parse_file_to_dict(f'{dirname}/rockstar.cfg')\n",
    "print(config_dict.keys())\n",
    "print(config_dict['Om'])\n",
    "print(type(config_dict['Om']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3480b9de-7dc7-4c1c-a5c3-9d3d9c8e45ca",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "borg_new",
   "language": "python",
   "name": "borg_new"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
