{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0 Set up your version of cave-client "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only run this one the first time to set up the token, then save your token in a file called token.txt\n",
    "\n",
    "import caveclient\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import time\n",
    "import csv\n",
    "import os\n",
    "\n",
    "\n",
    "client = caveclient.CAVEclient()\n",
    "client.auth.setup_token(make_new=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('token.txt', 'r') as file:\n",
    "    my_token = file.read().strip()\n",
    "\n",
    "datastack_name = \"flywire_fafb_public\"\n",
    "client = caveclient.CAVEclient(datastack_name, auth_token=my_token)\n",
    "\n",
    "# Print all the available versions and relative timestamps\n",
    "for version in client.materialize.get_versions():\n",
    "    print(f\"Version {version}: {client.materialize.get_timestamp(version)}\")\n",
    "\n",
    "# Print all the available tables for queries\n",
    "client.materialize.get_tables()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1 Fetch neurons from database\n",
    "## 1.1 By neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#this is to fetch neurons from the database\n",
    "\n",
    "cell_class_type_annos_df = client.materialize.query_table(\"hierarchical_neuron_annotations\", filter_in_dict={\"classification_system\": [\"cell_class\"]}) #cell type is another good one to look at\n",
    "cell_class_type_annos_df\n",
    "root_ids_cell_types = cell_class_type_annos_df[['pt_root_id', 'cell_type']]\n",
    "root_ids_cell_types.to_csv('classes_by_cell_type.csv', index=False)\n",
    "unique_cell_types = cell_class_type_annos_df['cell_type'].unique()\n",
    "print(unique_cell_types)\n",
    "\n",
    "def read_csv_data(csv_file, cell_types=None):\n",
    "    # FIX: column name was 'pt_root_id' instead of 'id', which was causing ValueError\n",
    "    data = []\n",
    "    with open(csv_file, 'r') as file:\n",
    "        csv_reader = csv.reader(file)\n",
    "        header = next(csv_reader)  # Skip the header row\n",
    "        id_index = header.index('pt_root_id')\n",
    "        cell_type_index = header.index('cell_type')\n",
    "        for row in csv_reader:\n",
    "            cell_id = int(row[id_index])\n",
    "            cell_type = row[cell_type_index]\n",
    "            if cell_types is None or cell_type in cell_types:\n",
    "                data.append((cell_id, cell_type))\n",
    "    return data\n",
    "\n",
    "csv_file = 'classes_by_cell_type.csv'\n",
    "cell_types = ['CX', 'Kenyon_Cell']  # Example list of cell types\n",
    "data = read_csv_data(csv_file, cell_types)\n",
    "print(data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 and by synapse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is to search from the database synapses\n",
    "\n",
    "\n",
    "#  Fetch the list of neurons\n",
    "neurons = client.materialize.query_table('proofread_neurons')\n",
    "\n",
    "# Extract unique neuron IDs\n",
    "neuron_ids = neurons['pt_root_id'].unique()\n",
    "\n",
    "# Create a mapping from neuron ID to matrix index\n",
    "neuron_id_to_index = {neuron_id: index for index, neuron_id in enumerate(neuron_ids)}\n",
    "\n",
    "if os.path.isfile('synapses.csv'):\n",
    "    already_logged = set(pd.read_csv('synapses.csv')['pre_pt_root_id'])\n",
    "else:\n",
    "    already_logged = set()\n",
    "\n",
    "# Open the CSV file in write mode\n",
    "with open('synapses.csv', 'a', newline='') as csvfile:\n",
    "    csv_writer = csv.writer(csvfile)\n",
    "    # Check if the file is empty\n",
    "    if csvfile.tell() == 0:\n",
    "        # Write the header row\n",
    "        csv_writer.writerow(['id', 'pre_pt_root_id', 'post_pt_root_id', 'connection_score', 'cleft_score', 'gaba', 'ach', 'glut', 'oct', 'ser', 'da', 'valid_nt', 'pre_pt_supervoxel_id', 'post_pt_supervoxel_id', 'neuropil', 'pre_pt_position', 'post_pt_position'])\n",
    "        \n",
    "    # Iterate through each neuron to fetch its synapses\n",
    "    for i, neuron_id in enumerate(neuron_ids):\n",
    "        # Check if the neuron ID is already in the first column of the CSV file\n",
    "        if neuron_id not in already_logged:\n",
    "            # Fetch the synapses for the neuron\n",
    "            neuron_start_time = time.time()\n",
    "            flag = 0\n",
    "            retries = 0\n",
    "            max_retries = 5\n",
    "            while flag == 0:\n",
    "                try:\n",
    "                    synapses = client.materialize.query_view(\"valid_synapses_nt_np_v6\", filter_in_dict={'pre_pt_root_id': [neuron_id]})\n",
    "                    flag = 1\n",
    "                except Exception as e:\n",
    "                    # FIX: bare except was swallowing all exceptions (including KeyboardInterrupt)\n",
    "                    # and never printing the cause. Now it logs the error and retries with a limit.\n",
    "                    retries += 1\n",
    "                    print(f\"Lost connection on neuron {i} ({e}), retrying ({retries}/{max_retries})...\")\n",
    "                    if retries >= max_retries:\n",
    "                        print(f\"Giving up on neuron {i} after {max_retries} retries.\")\n",
    "                        break\n",
    "                    time.sleep(10)\n",
    "                    with open('token.txt', 'r') as file:\n",
    "                        my_token = file.read().strip()\n",
    "\n",
    "                    datastack_name = \"flywire_fafb_public\"\n",
    "                    client = caveclient.CAVEclient(datastack_name, auth_token=my_token)\n",
    "\n",
    "            if flag == 0:\n",
    "                continue\n",
    "\n",
    "            for _, synapse in synapses.iterrows():\n",
    "                list_to_write = []\n",
    "                for elem in ['id', 'pre_pt_root_id', 'post_pt_root_id', 'connection_score', 'cleft_score', 'gaba', 'ach', 'glut', 'oct', 'ser', 'da', 'valid_nt', 'pre_pt_supervoxel_id', 'post_pt_supervoxel_id', 'neuropil', 'pre_pt_position', 'post_pt_position']:\n",
    "                    list_to_write += [synapse[elem]]    \n",
    "                \n",
    "                # Write the synapse details to the CSV file\n",
    "                csv_writer.writerow(list_to_write)\n",
    "            \n",
    "            # Calculate elapsed time and estimate remaining time\n",
    "            elapsed_time = time.time() - neuron_start_time\n",
    "            print(f\"Processed neuron {i+1}/{len(neuron_ids)} in {elapsed_time:.2f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 This is just to check the presence of specific cells inside a .csv file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#this it too check stuff from a specific .csv file\n",
    "import csv\n",
    "\n",
    "csv_file = 'classes_by_cell_type.csv'\n",
    "cell_types = ['olfactory', 'visual', 'mechanosensory', 'hygrosensory', 'unknown_sensory', 'ocellar', 'gustatory', 'thermosensory']  # Example list of cell types\n",
    "\n",
    "data = []\n",
    "with open(csv_file, 'r') as file:\n",
    "    csv_reader = csv.reader(file)\n",
    "    header = next(csv_reader)  # Skip the header row\n",
    "    id_index = header.index('pt_root_id')\n",
    "    cell_type_index = header.index('cell_type')\n",
    "    for row in csv_reader:\n",
    "        cell_id = int(row[id_index])\n",
    "        cell_type = row[cell_type_index]\n",
    "        if cell_types is None or cell_type in cell_types:\n",
    "            data.append((cell_id, cell_type))\n",
    "\n",
    "data_dict = {cell_id: cell_type for cell_id, cell_type in data}\n",
    "\n",
    "\n",
    "\n",
    "print(data)\n",
    "len(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 This creates a graph from the .csv (make sure you have the .csv in the right place). IMPORTANT: threshold_connection and threshold_cleft are your quality selection parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import csv\n",
    "import networkx as nx\n",
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "import pickle\n",
    "import os\n",
    "\n",
    "threshold_connection = 100 # at the moment this is fully arbitrary and should be changed\n",
    "threshold_cleft = 50\n",
    "\n",
    "# here is the function that will determine the sign of the neurotransmitter\n",
    "\n",
    "def synapse_characteristics(cell_type, neurotransmitter):\n",
    "    # General neurotransmitter effects for excitatory/inhibitory/modulatory\n",
    "    neurotransmitter_effects = {\n",
    "        \"gaba\": \"inhibitory\",\n",
    "        \"ach\": \"excitatory\",\n",
    "        \"glut\": \"excitatory\",\n",
    "        \"oct\": \"modulatory\",\n",
    "        \"ser\": \"modulatory\",\n",
    "        \"da\": \"modulatory\"\n",
    "    }\n",
    "\n",
    "    # Cell type-specific effects and threshold modulation patterns\n",
    "    cell_type_effects = {\n",
    "        \"kenyon_cell\": {\"ach\": (\"excitatory\", \"lower\")},\n",
    "        \"alpn\": {\"glut\": (\"excitatory\", None)},\n",
    "        \"alln\": {\"gaba\": (\"inhibitory\", None), \"ach\": (\"mixed\", None)},\n",
    "        \"dan\": {\"da\": (\"modulatory\", \"lower\")},\n",
    "        \"mbon\": {\"da\": (\"modulatory\", \"lower\"), \"ach\": (\"excitatory\", None), \"gaba\": (\"inhibitory\", None)},\n",
    "        \"lo\": {\"glut\": (\"excitatory\", None), \"ach\": (\"excitatory\", None)},\n",
    "        \"me\": {\"glut\": (\"excitatory\", None), \"ach\": (\"excitatory\", None)},\n",
    "        \"visual\": {\"glut\": (\"excitatory\", None)},\n",
    "        \"olfactory\": {\"glut\": (\"excitatory\", None)},\n",
    "        \"mechanosensory\": {\"glut\": (\"excitatory\", \"lower\"), \"oct\": (\"modulatory\", \"lower\")},\n",
    "        \"clock\": {\"da\": (\"modulatory\", \"lower\"), \"ser\": (\"modulatory\", \"raise\")},\n",
    "        \"lhcent\": {\"gaba\": (\"inhibitory\", None), \"glut\": (\"excitatory\", None), \"ser\": (\"modulatory\", \"raise\")},\n",
    "        \"alin\": {\"gaba\": (\"inhibitory\", None)},\n",
    "        \"lhln\": {\"gaba\": (\"inhibitory\", None)},\n",
    "        \"tpn\": {\"oct\": (\"modulatory\", \"lower\")}\n",
    "    }\n",
    "\n",
    "    # Standardize inputs to lowercase for matching\n",
    "    neurotransmitter = neurotransmitter.lower()\n",
    "    cell_type = cell_type.lower()\n",
    "\n",
    "    # Check if the cell type has specific rules for the neurotransmitter\n",
    "    if cell_type in cell_type_effects:\n",
    "        cell_rules = cell_type_effects[cell_type]\n",
    "        if neurotransmitter in cell_rules:\n",
    "            effect, threshold_modulation = cell_rules[neurotransmitter]\n",
    "            return effect, threshold_modulation\n",
    "\n",
    "    # Default to general neurotransmitter effects if no cell-specific rule found\n",
    "    effect = neurotransmitter_effects.get(neurotransmitter, \"unknown\")\n",
    "    threshold_modulation = None  # Default to None if no specific modulation information\n",
    "    return effect, threshold_modulation\n",
    "\n",
    "# Step 0: get the class of every cell\n",
    "\n",
    "csv_file = 'classes_by_cell_type.csv'\n",
    "cell_classes = {}\n",
    "biases = {}\n",
    "with open(csv_file, 'r') as file:\n",
    "    csv_reader = csv.reader(file)\n",
    "    header = next(csv_reader)  # Skip the header row\n",
    "    id_index = header.index('pt_root_id')\n",
    "    cell_type_index = header.index('cell_type')\n",
    "    for row in csv_reader:\n",
    "        cell_id = row[id_index]\n",
    "        cell_type = row[cell_type_index]\n",
    "        if cell_id not in cell_classes.keys():\n",
    "            cell_classes.update({cell_id: cell_type})\n",
    "            biases.update({cell_id: 0})\n",
    "\n",
    "neuro_transmitters = [\"gaba\",\"ach\",\"glut\",\"oct\",\"ser\",\"da\"]\n",
    "\n",
    "# Step 1: Read the CSV file\n",
    "# FIX: file name was 'synapses_complete.csv' but it was never created anywhere.\n",
    "# Since the file written in 1.2 is 'synapses.csv', it should be used here as well.\n",
    "with open('synapses.csv', 'r') as file:\n",
    "    print('initializing reader')\n",
    "    csv_reader = csv.reader(file)\n",
    "    print('reader done')\n",
    "    header = next(csv_reader)  # Skip the header row\n",
    "\n",
    "    # Step 2: Initialize a NetworkX graph\n",
    "    print('initializing graph')\n",
    "    G = nx.DiGraph()\n",
    "    print('graph done')\n",
    "\n",
    "    # Step 3: Add edges to the graph\n",
    "    current_row = 0\n",
    "\n",
    "    # FIX: file was already opened and header skipped; file.seek(0) + next()\n",
    "    # was doing the same thing redundantly. Removed.\n",
    "\n",
    "    # Iterate through each row in the CSV file\n",
    "    unique_neurons = 0\n",
    "    for row in csv_reader:\n",
    "        if float(row[3]) > threshold_connection and float(row[4]) > threshold_cleft and row[1] in cell_classes.keys() and row[2] in cell_classes.keys():\n",
    "            source = row[1]  # Assuming the source column is at index 1\n",
    "            target = row[2]  # Assuming the target column is at index 2\n",
    "            weight = float(row[3])  # Assuming the weight column is at index 3\n",
    "            neuro_trs = neuro_transmitters[np.argmax([float(row[5]), float(row[6]), float(row[7]), float(row[8]), float(row[9]), float(row[10])])]\n",
    "\n",
    "            if target in cell_classes.keys():\n",
    "                tmp_cell_class = cell_classes[target]\n",
    "            else:\n",
    "                tmp_cell_class = \"unknown\"\n",
    "            effect, threshold_modulation = synapse_characteristics(tmp_cell_class, neuro_trs)\n",
    "\n",
    "            if effect == \"mixed\":\n",
    "                effect = random.choice([\"excitatory\", \"inhibitory\"])\n",
    "\n",
    "            # FIX: target was already added to G as string, but the comparison checked 'int(target) not in G.nodes()'\n",
    "            # -> due to type mismatch the condition was always True.\n",
    "            is_new_node = target not in G.nodes()\n",
    "\n",
    "            if effect == \"excitatory\":\n",
    "                G.add_edge(source, target, weight=weight)\n",
    "            elif effect == \"inhibitory\":\n",
    "                G.add_edge(source, target, weight=-weight)\n",
    "            elif effect == \"modulatory\":\n",
    "                if threshold_modulation == \"lower\": # there will probably be the need of modulation for these\n",
    "                    biases[target] -= weight\n",
    "                elif threshold_modulation == \"raise\":\n",
    "                    biases[target] += weight\n",
    "\n",
    "            if effect in (\"excitatory\", \"inhibitory\") and is_new_node:\n",
    "                unique_neurons += 1\n",
    "\n",
    "        # Update loading bar\n",
    "        current_row += 1\n",
    "        if current_row % 100000 == 0:\n",
    "            print(f\"\\rRows processed: {current_row}     Unique neurons found: {unique_neurons} \", end='', flush=True)\n",
    "    print(\"\\nProcessing complete.\")\n",
    "\n",
    "    directory = f\"networks_graphs/Con{threshold_connection}_Cleft{threshold_cleft}\"\n",
    "    # FIX: added exist_ok=True to prevent FileExistsError if the directory already exists\n",
    "    os.makedirs(directory, exist_ok=True)\n",
    "\n",
    "# Save the graph\n",
    "is_directed = nx.is_directed(G)\n",
    "print(f\"Is the graph directed? {is_directed}\")\n",
    "nx.write_graphml(G, f\"{directory}/graph.graphml\")\n",
    "# Save biases to a .pkl file\n",
    "with open(f\"{directory}/biases.pkl\", 'wb') as file:\n",
    "    pickle.dump(biases, file)\n",
    "\n",
    "if False:\n",
    "\n",
    "    #Step 4: Extract edge list\n",
    "    edge_index = torch.tensor([(int(u), int(v)) for u, v in G.edges]).t().contiguous()\n",
    "    edge_weight = torch.tensor([G[u][v]['weight'] for u, v in G.edges])\n",
    "\n",
    "    # Step 5: Create a PyTorch sparse matrix\n",
    "    sparse_matrix = torch.sparse_coo_tensor(edge_index, edge_weight, (G.number_of_nodes(), G.number_of_nodes()))\n",
    "\n",
    "    print(sparse_matrix)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "\n",
    "# Define the directory and file path\n",
    "directory = f\"networks_graphs/Con{threshold_connection}_Cleft{threshold_cleft}\"\n",
    "graph_path = f\"{directory}/graph.graphml\"\n",
    "\n",
    "# Load the graph\n",
    "G = nx.read_graphml(graph_path)\n",
    "if nx.is_directed(G):\n",
    "    print(\"The graph is directed.\")\n",
    "else:\n",
    "    print(\"The graph is undirected.\")\n",
    "\n",
    "print(f\"Graph loaded with {G.number_of_nodes()} nodes and {G.number_of_edges()} edges.\")\n",
    "\n",
    "# Keep only the largest connected component\n",
    "largest_cc = max(nx.weakly_connected_components(G), key=len)\n",
    "G = G.subgraph(largest_cc).copy()\n",
    "\n",
    "print(f\"Graph reduced to largest connected component with {G.number_of_nodes()} nodes and {G.number_of_edges()} edges.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "# Convert the graph G to a scipy sparse matrix\n",
    "sparse_matrix = nx.to_scipy_sparse_array(G, weight='weight', format='csr')\n",
    "\n",
    "# Compute the sparsity of the matrix\n",
    "total_elements = sparse_matrix.shape[0] * sparse_matrix.shape[1]\n",
    "nonzero_elements = sparse_matrix.nnz\n",
    "sparsity = 1 - (nonzero_elements / total_elements)\n",
    "print(f\"Sparsity of the matrix: {sparsity:.4f}\")\n",
    "\n",
    "# Check if the matrix is symmetrical\n",
    "is_symmetric = (sparse_matrix != sparse_matrix.T).nnz == 0\n",
    "print(f\"Is the matrix symmetrical? {'Yes' if is_symmetric else 'No'}\")\n",
    "\n",
    "# Save the sparse matrix to a pickle file\n",
    "with open('W_whole_world.pkl', 'wb') as f:\n",
    "    pickle.dump(sparse_matrix, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fly_connectome_2",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}