WS-KCNHUB/machines/kcnhub/servers/jupyterhub.nix

{ config, pkgs, lib, ... }: let
  aiomcache = pkgs.python3.pkgs.buildPythonPackage rec {
    pname  = "aio-memcache";
    version = "0.8.1";

    src = pkgs.fetchFromGitHub {
      owner = "aio-libs";
      repo = "aiomcache";
      rev = "v${version}";
      hash = "sha256-oRMN1seEjFSsq4wjkIXHM7Osq8y/5WFExGCEr6eM9vc=";
    };

    propagatedBuildInputs = with pkgs.python3.pkgs; [
      python-memcached
    ];

    # Networking required
    doCheck = false;
  };
  aiohttp-session = pkgs.python3.pkgs.buildPythonPackage rec {
    pname = "aiohttp-session";
    version = "2.12.0";

    src = pkgs.fetchFromGitHub {
      owner = "aio-libs";
      repo = "aiohttp-session";
      rev = "v${version}";
      hash = "sha256-7MNah4OIQnoxLoZkLOdeu5uCwSyPMhc6Wsht8dFconc=";
    };

    propagatedBuildInputs = with pkgs.python3.pkgs; [
      aiohttp
      aiomcache
      cryptography
      multidict
      pynacl
      redis
      yarl
    ];

    # Checks require networking
    doCheck = false;
  };
  matlab-proxy = pkgs.python3.pkgs.buildPythonPackage rec {
    pname = "matlab_proxy";
    version = "0.12.1";
    format = "wheel";
    src = pkgs.python3.pkgs.fetchPypi {
      inherit pname version format;
      dist = "py3";
      python = "py3";
      hash = "sha256-TDZUJEXJYF5Vh2BG5o13H3wdh/WdVdoESXqL89ckc6I=";
    };
    
    propagatedBuildInputs = with pkgs.python3.pkgs; [
      aiohttp
      aiohttp-session
      psutil

      # For Xvfb, for plots and graphics
      pkgs.xorg.xorgserver
    ];
  };
  simpervisor = pkgs.python3.pkgs.buildPythonPackage rec {
    pname = "simpervisor";
    version = "1.0.0";
    format = "pyproject";

    src = pkgs.fetchFromGitHub {
      owner = "jupyterhub";
      repo = "simpervisor";
      rev = "v${version}";
      hash = "sha256-73vkiQtOT0M9Vww1nYZ76JR2koWt/NPIav46k1fHOzc=";
    };

    nativeBuildInputs = with pkgs.python3.pkgs; [
      hatchling
    ];
  };
  jupyter-server-proxy = pkgs.python3.pkgs.buildPythonPackage rec {
    pname = "jupyter_server_proxy";
    version = "4.1.0";
    format = "pyproject";

    src = pkgs.python3.pkgs.fetchPypi {
      inherit pname version;
      hash = "sha256-LPrDtCMv5xROjmApa0+GFwi08TspJgos8ol2v45hf3A=";
    };

    nativeBuildInputs = with pkgs.python3.pkgs; [
      hatchling
      hatch-jupyter-builder
      hatch-nodejs-version
      jupyterlab
    ];
    
    propagatedBuildInputs = with pkgs.python3.pkgs; [
      notebook
      simpervisor
      aiohttp
    ];
    
    # tests require network
    doCheck = false;
  };
  
  jupyter-matlab-proxy = pkgs.python3.pkgs.buildPythonPackage rec {
    pname = "jupyter_matlab_proxy";
    version = "0.11.0";
    format = "wheel";

    src = pkgs.python3.pkgs.fetchPypi {
      inherit pname version format;
      python = "py3";
      dist = "py3";
      hash = "sha256-KzQgs1zXGPjONCnLvCNQkQBkkkn9J29YFi8BHbh+D1I=";
    };

    propagatedBuildInputs = with pkgs.python3.pkgs; [
      jupyter-server-proxy
      simpervisor
      matlab-proxy
      psutil
      requests
      ipykernel
    ];
  };
  in {
  services.jupyterhub = {
    enable = true;
    port = 6501;
    # https://github.com/jupyterhub/systemdspawner
    # Configuration options defined ^
    # mem_limit must be higher to allow for MATLAB
    extraConfig = ''
      c.SystemdSpawner.mem_limit = '8G'
      c.SystemdSpawner.cpu_limit = 2.0
      c.SystemdSpawner.isolate_devices = False
      c.SystemdSpawner.extra_paths = [
        '${pkgs.buildEnv {
          name = "jupyterhub-system-spawner-env";
          paths = with pkgs; [
            gcc
            # For MATLAB Kernel
            matlab
            # For matlab-proxy-app
            (pkgs.python3.withPackages (p: with p; [
              # For pkg_resources
              setuptools
              matlab-proxy
              jupyter-matlab-proxy
            ]))
          ];
        }}/bin']
      c.SystemdSpawner.environment["MATLAB_INSTALL_DIR"] = "${config.environment.sessionVariables.MATLAB_INSTALL_DIR}"
      c.SystemdSpawner.environment["LM_LICENSE_FILE"] = "${config.environment.sessionVariables.LM_LICENSE_FILE}"
      c.SystemdSpawner.environment["MLM_LICENSE_FILE"] = "${config.environment.sessionVariables.MLM_LICENSE_FILE}"
      # c.SystemdSpawner.environment["MWI_JUPYTER_LOG_LEVEL"] = "DEBUG"
    '';
    # Extra Paths used to add gcc et al for packages like Brian2
    # https://github.com/jupyterhub/systemdspawner#extra_paths

    jupyterlabEnv = pkgs.python3.withPackages (p: with p; [
      jupyterhub
      jupyterlab
      # Extensions:
      jupyter-collaboration
      jupyter-server-mathjax
      jupyter-console
      jupyterlab-lsp
      python-lsp-server
      # MATLAB Proxy Needs to be here in order to modify and add the "launch matlab" etc.
      jupyter-matlab-proxy # <-- This in effect adds the kernel
    ]);
    
    kernels = {
      python3 = let
        env = (pkgs.python3.withPackages (pythonPackages: with pythonPackages; [
          # Necessary for use as a kernel
          ipykernel
          ipdb
          # Interactive widgets
          ipywidgets
          # Show images/media
          mediapy
          # Progress bars etc. for in Jupyter/IPython
          halo

          # Common Python Libraries
          numpy
          scipy
          scikit-learn
          pandas
          # Tables displayed in Jupyter: like Pandas Dataframes
          ipytablewidgets
          matplotlib
          ipympl
          matplotlib-inline
          seaborn

          # Neural Data Processing Libraries
          # Temporarily broken?: https://github.com/NixOS/nixpkgs/issues/259812
          # mne-python
          nibabel
          # View niftis!
          ipyniivue
          neo
          
          # Neurosimulators
          neuronpy
          # nest with MPI is incompatible with Jupyter:
          # https://www.nest-simulator.org/pynest-api/_modules/nest.html
          # see pynestkernel comments about their workaround
          (nest.override {
            withMpi = false;
          })
          brian2

          # Machine learning toolkits:
          tensorflow
          tensorboard
          keras
          edward
          transformers
          tflearn
          torch
          torchvision
          torchsde
          torchmetrics
          torchio
          torchdiffeq
          botorch
          lion-pytorch
        ])).override (args: { ignoreCollisions = true; });
        # Odd collision between tensorboard of torch & tensorflow
        # need to resolve later
      in {
        displayName = "Python 3 for Computational Neuroscience";
        argv = [
          "${env.interpreter}"
          "-m"
          "ipykernel_launcher"
          "-f"
          "{connection_file}"
        ];
        language = "python";
        logo32 = "${env}/${env.sitePackages}/ipykernel/resources/logo-32x32.png";
        logo64 = "${env}/${env.sitePackages}/ipykernel/resources/logo-64x64.png";
      };
    };
  };
  
  services.caddy.virtualHosts = {
    "jupyter.ws.kcnhub.com" = {
      extraConfig = ''
        reverse_proxy 127.0.0.1:${toString config.services.jupyterhub.port}
      '';           
    };
  };
}