{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "61421059-d31e-4fb9-a817-a930a332955d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zarr\n",
    "import dask.array as da\n",
    "from tqdm import tqdm\n",
    "import yaml\n",
    "#from numcodecs import Blosc #hua add to replace zarr.Blosc\n",
    "\n",
    "def merge_cap_predictions(input_dir, output_path, celltype, cap_list):\n",
    "    input_dir = os.path.join(input_dir, celltype)\n",
    "    output_path = os.path.join(output_path, celltype)\n",
    "    os.makedirs(output_path, exist_ok=True)\n",
    "\n",
    "    cap_list = sorted(cap_list)\n",
    "\n",
    "    # Save cap list to file\n",
    "    with open(os.path.join(output_path, f'cap_list.txt'), 'w') as f:\n",
    "        for cap in cap_list:\n",
    "            f.write(f\"{cap}\\n\")\n",
    "\n",
    "    # Construct paths to zarr files for each CAP\n",
    "    zarr_files = []\n",
    "    for cap in cap_list:\n",
    "        zarr_path = os.path.join(input_dir, cap, 'processed', 'data.zarr')\n",
    "        if os.path.exists(zarr_path):\n",
    "            zarr_files.append(zarr_path)\n",
    "        else:\n",
    "            print(f\"Warning: {zarr_path} does not exist, skipping\")\n",
    "    \n",
    "    if not zarr_files:\n",
    "        raise ValueError(\"No valid zarr files found to merge\")\n",
    "\n",
    "    # Get chrs from first zarr file\n",
    "    infer_chrs = zarr.open(zarr_files[0], mode='r')['chrs']\n",
    "    infer_chrs = [c for c in infer_chrs]\n",
    "\n",
    "    # 3. Open or create output Zarr store (directory store on disk in this case)\n",
    "    output_file = os.path.join(output_path, f'data.zarr')\n",
    "    root = zarr.group(store = output_file, overwrite = False) # init zarr group\n",
    "\n",
    "    # 4. Iterate over each chromosome to process\n",
    "    \n",
    "    for chrom in tqdm(infer_chrs, desc=\"Processing chromosomes\"):\n",
    "        files = [os.path.join(input_dir, cap, 'processed', 'data.zarr', f'chrs/{chrom}') for cap in cap_list]\n",
    "        # Load each 1D Zarr file as a Dask array (lazy loaded, no data read yet)\n",
    "        arrays = []\n",
    "        for i, f in enumerate(tqdm(files, desc=f\"Loading arrays for {chrom}\", leave=False)):\n",
    "            arrays.append(da.from_zarr(f))\n",
    "        # (All arrays must have the same length for stacking; ensure shapes align)\n",
    "        \n",
    "        # 5. Stack arrays along a new axis (axis=0 creates shape (n_files, length))\n",
    "        print(f\"Stacking {len(arrays)} arrays for {chrom}...\")\n",
    "        # Stack arrays along the second axis (axis=1) instead of the first\n",
    "        # This will create a shape of (length, n_files) instead of (n_files, length)\n",
    "        stacked = da.stack(arrays, axis=1)\n",
    "        \n",
    "        # 6. Rechunk to have chunk size appropriate for the new shape (1_000_000, n)\n",
    "        n_files = len(files)\n",
    "        cap_chunk_size = min(n_files, 100)\n",
    "        chunk_size = (1_000_000, cap_chunk_size)\n",
    "        print(f\"Rechunking array for {chrom} with chunk size {chunk_size}...\")\n",
    "        stacked = stacked.rechunk(chunks=chunk_size)\n",
    "        \n",
    "        # 7. Write this stacked array to the Zarr store under chrs/<chrom>\n",
    "        # Use component to specify the subgroup path, and allow overwrite if exists.\n",
    "        print(f\"Writing {chrom} to zarr store...\")\n",
    "        da.to_zarr(stacked, output_file, component=f\"chrs/{chrom}\", overwrite=True, compressor=zarr.Blosc(cname='zstd', clevel=3, shuffle=zarr.Blosc.SHUFFLE))\n",
    "\n",
    "        print(f\"Completed {chrom}: shape {stacked.shape}, chunk size {stacked.chunksize}\")\n",
    "\n",
    "def load_inputs(config):\n",
    "    celltype_list = read_list(os.path.join(config['inference_config']['input']['root'], config['inference_config']['input']['celltype_list_path']))\n",
    "    cap_list = read_list(os.path.join(config['inference_config']['input']['root'], config['inference_config']['input']['cap_list_path']))\n",
    "    return celltype_list, cap_list\n",
    "\n",
    "def read_list(list_path):\n",
    "    item_list = []\n",
    "    with open(list_path, 'r') as file:\n",
    "        for line in file:\n",
    "            item_list.append(line.strip())\n",
    "    return item_list\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a742e67b-425c-45ba-9964-0879c9a8064a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    # Load config\n",
    "    with open('config.yaml', 'r') as f:\n",
    "        config = yaml.safe_load(f)\n",
    "\n",
    "    input_dir = config['inference_config']['output']['path']\n",
    "    output_path = config['merge_inference_config']['output']['path']\n",
    "\n",
    "    celltype_list, cap_list = load_inputs(config)\n",
    "\n",
    "    print(celltype_list)\n",
    "\n",
    "    for celltype in celltype_list:\n",
    "        merge_cap_predictions(input_dir, output_path, celltype, cap_list)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}