{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f48b9a60",
   "metadata": {},
   "source": [
    "# Solutions to Notebook Exercises\n",
    "\n",
    "## Julia Basics: Exercise 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a06fd02a",
   "metadata": {},
   "outputs": [],
   "source": [
    "function ex1(a)\n",
    "    j = 1\n",
    "    m = a[j]\n",
    "    for (i,ai) in enumerate(a)\n",
    "        if m < ai\n",
    "            m = ai\n",
    "            j = i\n",
    "        end\n",
    "    end\n",
    "    (m,j)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "175b6c35",
   "metadata": {},
   "source": [
    "## Julia Basics: Exercise 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb289acd",
   "metadata": {},
   "outputs": [],
   "source": [
    "ex2(f,g) = x -> f(x) + g(x) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86250e27",
   "metadata": {},
   "source": [
    "## Julia Basics: Exercise 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41b537ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "function compute_values(n,max_iters)\n",
    "    x = LinRange(-1.7,0.7,n)\n",
    "    y = LinRange(-1.2,1.2,n)\n",
    "    values = zeros(Int,n,n)\n",
    "    for j in 1:n\n",
    "        for i in 1:n\n",
    "            values[i,j] = mandel(x[i],y[j],max_iters)\n",
    "        end\n",
    "    end\n",
    "    values\n",
    "end\n",
    "values = compute_values(1000,10)\n",
    "using GLMakie\n",
    "heatmap(x,y,values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6d12733",
   "metadata": {},
   "source": [
    "## Matrix Multiplication : Exercise 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be73e87a",
   "metadata": {},
   "outputs": [],
   "source": [
    "function matmul_dist_3!(C,A,B)\n",
    "    m = size(C,1)\n",
    "    n = size(C,2)\n",
    "    l = size(A,2)\n",
    "    @assert size(A,1) == m\n",
    "    @assert size(B,2) == n\n",
    "    @assert size(B,1) == l\n",
    "    @assert mod(m,nworkers()) == 0\n",
    "    # Implement here\n",
    "    nrows_w = div(m,nworkers())\n",
    "    @sync for (i,w) in enumerate(workers())\n",
    "        rows_w = (1:nrows_w) .+ (i-1)*nrows_w\n",
    "        Aw = A[rows_w,:]\n",
    "        ftr = @spawnat w begin\n",
    "            Cw = similar(Aw,nrows_w,n)\n",
    "            matmul_seq!(Cw,Aw,B)\n",
    "            Cw\n",
    "        end\n",
    "        @async C[rows_w,:] = fetch(ftr)\n",
    "    end\n",
    "    C\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d9f4813",
   "metadata": {},
   "source": [
    "## Jacobi Method : Exercise 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf3b1e72",
   "metadata": {},
   "outputs": [],
   "source": [
    "@mpi_do manager begin\n",
    "    using MPI\n",
    "    comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "    nw = MPI.Comm_size(comm)\n",
    "    iw = MPI.Comm_rank(comm)+1\n",
    "    function jacobi_mpi(n,niters)\n",
    "        if mod(n,nw) != 0\n",
    "            println(\"n must be a multiple of nw\")\n",
    "            MPI.Abort(comm,1)\n",
    "        end\n",
    "        n_own = div(n,nw)\n",
    "        u = zeros(n_own+2)\n",
    "        u[1] = -1\n",
    "        u[end] = 1\n",
    "        u_new = copy(u)\n",
    "        for t in 1:niters\n",
    "            reqs_snd = MPI.Request[]\n",
    "            reqs_rcv = MPI.Request[]\n",
    "            if iw != 1\n",
    "                neig_rank = (iw-1)-1\n",
    "                req = MPI.Isend(view(u,2:2),comm,dest=neig_rank,tag=0)\n",
    "                push!(reqs_snd,req)\n",
    "                req = MPI.Irecv!(view(u,1:1),comm,source=neig_rank,tag=0)\n",
    "                push!(reqs_rcv,req)\n",
    "            end\n",
    "            if iw != nw\n",
    "                neig_rank = (iw+1)-1\n",
    "                s = n_own+1\n",
    "                r = n_own+2\n",
    "                req = MPI.Isend(view(u,s:s),comm,dest=neig_rank,tag=0)\n",
    "                push!(reqs_snd,req)\n",
    "                req = MPI.Irecv!(view(u,r:r),comm,source=neig_rank,tag=0)\n",
    "                push!(reqs_rcv,req)\n",
    "            end\n",
    "            for i in 3:n_own\n",
    "                u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "            end\n",
    "            MPI.Waitall(reqs_rcv)\n",
    "            for i in (2,n_own+1)\n",
    "                u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "            end\n",
    "            MPI.Waitall(reqs_snd)\n",
    "            u, u_new = u_new, u\n",
    "        end\n",
    "        u\n",
    "        @show u\n",
    "    end\n",
    "    niters = 100\n",
    "    load = 4\n",
    "    n = load*nw\n",
    "    jacobi_mpi(n,niters)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f343157",
   "metadata": {},
   "source": [
    "## Exercise: Ring communication - MPI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a49be691",
   "metadata": {},
   "outputs": [],
   "source": [
    "using MPI\n",
    "using Test\n",
    "\n",
    "MPI.Init()\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = MPI.Comm_rank(comm)\n",
    "id = rank + 1\n",
    "root = 0\n",
    "size = MPI.Comm_size(comm)\n",
    "\n",
    "dst = mod(rank + 1, size)\n",
    "src = mod(rank - 1, size)\n",
    "\n",
    "send_buf = id\n",
    "recv_buf = 1\n",
    "\n",
    "if rank == root \n",
    "    # Proc 1: Send id async to destination, then wait for receive\n",
    "    MPI.isend(send_buf, comm; dest=dst, tag=0)\n",
    "    recv_buf = MPI.recv(comm; source=src, tag=0)\n",
    "    @show recv_buf == factorial(size)\n",
    "    @test recv_buf == factorial(size)\n",
    "else\n",
    "    # Other procs: receive sync and send async to next process\n",
    "    recv_buf = MPI.recv(comm; source=src, tag=0)\n",
    "    send_buf = recv_buf * id\n",
    "    MPI.isend(send_buf, comm; dest=dst, tag=0)\n",
    "end\n",
    "\n",
    "MPI.Finalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cbbf074",
   "metadata": {},
   "source": [
    "## Exercise: Ring communication - Distributed.jl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc156523",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Distributed \n",
    "using Test\n",
    "\n",
    "np = 4\n",
    "add_n = np - nprocs() \n",
    "addprocs(add_n)\n",
    "worker_ids = workers()\n",
    "@assert nprocs() > nworkers()\n",
    "\n",
    "# Initialize id channel\n",
    "id_chnl = RemoteChannel(()->Channel{Int}(1))\n",
    "put!(id_chnl, 1)\n",
    "\n",
    "# Initialize data channel\n",
    "job_chnl = RemoteChannel(()->Channel{Int}(1))\n",
    "put!(job_chnl, 1)\n",
    "\n",
    "@sync for w in workers()\n",
    "    @spawnat w begin\n",
    "        pos = findfirst(worker_ids .== w) + 1\n",
    "        dst = mod(pos, np) + 1\n",
    "        src = mod(pos-2, np) + 1\n",
    "        while true \n",
    "            pred = fetch(id_chnl)\n",
    "            if pred == src\n",
    "                take!(id_chnl)\n",
    "                value = take!(job_chnl)\n",
    "                put!(job_chnl, value * pos)  \n",
    "                put!(id_chnl, pos)  \n",
    "                break\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "res = take!(job_chnl)\n",
    "@show res\n",
    "@test res == factorial(np)\n",
    "\n",
    "rmprocs(workers())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "508125f9",
   "metadata": {},
   "source": [
    "### New ring example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af4cfd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere workers() begin\n",
    "    using MPI\n",
    "    MPI.Init()\n",
    "    comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    buffer = Ref(0)\n",
    "    if rank == 0\n",
    "        msg = rand(1:10)\n",
    "        buffer[] = msg\n",
    "        println(\"msg = $(buffer[])\")\n",
    "        MPI.Send(buffer,comm;dest=rank+1,tag=0)\n",
    "        MPI.Recv!(buffer,comm;source=nranks-1,tag=0)\n",
    "        println(\"msg = $(buffer[])\")\n",
    "    else\n",
    "        dest = (rank != nranks-1) ? rank+1 : 0\n",
    "        MPI.Recv!(buffer,comm;source=rank-1,tag=0)\n",
    "        buffer[] += 1\n",
    "        println(\"msg = $(buffer[])\")\n",
    "        MPI.Send(buffer,comm;dest,tag=0)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30ece423",
   "metadata": {},
   "outputs": [],
   "source": [
    "ftrs = Vector{Future}(undef,nprocs())\n",
    "for p in 1:nprocs()\n",
    "    ftrs[p] = @spawnat p RemoteChannel(()->Channel{Int}(1))\n",
    "end\n",
    "for p in 1:nprocs()\n",
    "    @spawnat p begin        \n",
    "        chnl_snd = fetch(ftrs[p])\n",
    "        source = (p != 1) ? p-1 : nprocs()\n",
    "        chnl_rcv = fetch(ftrs[source])\n",
    "        if p == 1\n",
    "            msg = rand(1:10)\n",
    "            @show msg\n",
    "            put!(chnl_snd,msg)\n",
    "            msg = take!(chnl_rcv)\n",
    "            @show msg\n",
    "        else\n",
    "            msg = take!(chnl_rcv)\n",
    "            msg += 1\n",
    "            @show msg\n",
    "            put!(chnl_snd,msg)\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "761ce452",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function work(msg)\n",
    "    msg += 1\n",
    "    println(\"msg = $msg\")\n",
    "    if myid() == nprocs()\n",
    "        @spawnat 1 println(\"msg = $msg\")\n",
    "    else\n",
    "        next = myid() + 1\n",
    "        @spawnat next work(msg)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48e594e7",
   "metadata": {},
   "source": [
    "## LEQ: Exercise 1\n",
    "### Blockwise partitioning: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d890006",
   "metadata": {},
   "outputs": [],
   "source": [
    "function ge_par_block!(B,n,m,comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    # Init buffers\n",
    "    T = eltype(B)\n",
    "    if rank == 0\n",
    "        buffer_root = Vector{T}(undef,n*m)\n",
    "        buffer_root[:] = transpose(B)[:]\n",
    "    else\n",
    "        buffer_root = Vector{T}(undef,0)\n",
    "    end    \n",
    "    nw = div(n,nranks)\n",
    "    buffer =  Vector{T}(undef,nw*m)\n",
    "    # Send local matrices to workers\n",
    "    MPI.Scatter!(buffer_root,buffer,comm;root=0)\n",
    "    Bw = Matrix{T}(undef,nw,m)\n",
    "    transpose(Bw)[:] = buffer\n",
    "    MPI.Barrier(comm)\n",
    "    # time calcultation\n",
    "    t = @elapsed ge_block_worker!(Bw,n,comm)\n",
    "    # Gather results\n",
    "    buffer[:] = transpose(Bw)[:]\n",
    "    MPI.Gather!(buffer,buffer_root,comm;root=0)\n",
    "    Tf = typeof(t)\n",
    "    if rank == 0\n",
    "        transpose(B)[:] = buffer_root[:]\n",
    "        ts = Vector{Tf}(undef,nranks)\n",
    "    else\n",
    "        ts = Vector{Tf}(undef,0)\n",
    "    end\n",
    "    MPI.Gather!([t],ts,comm;root=0)\n",
    "    B, ts\n",
    "end \n",
    "\n",
    "function ge_block_worker!(Bw,n,comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    nw,m = size(Bw)\n",
    "    B_k = similar(Bw,m)\n",
    "    lb = nw*rank + 1\n",
    "    ub = nw*(rank+1)\n",
    "    @inbounds for k in 1:n\n",
    "        # If I have row k\n",
    "        if k in lb:ub\n",
    "            myk = k - lb + 1\n",
    "            # Update just row k\n",
    "            #@show Bw[myk,k]\n",
    "            for t in (k+1):m\n",
    "                Bw[myk,t] =  Bw[myk,t]/Bw[myk,k]\n",
    "            end\n",
    "            Bw[myk,k] = 1\n",
    "            # Send k to other procs\n",
    "            B_k .= view(Bw,myk,:)\n",
    "            MPI.Bcast!(B_k,comm;root=rank)\n",
    "        else\n",
    "            myk = 0\n",
    "            owner = div(k-1,nw)\n",
    "            MPI.Bcast!(B_k,comm;root=owner)\n",
    "        end\n",
    "        # Do nothing if I only have rows < k \n",
    "        if k <= ub\n",
    "            for i in (myk+1):nw\n",
    "                for j in (k+1):m\n",
    "                    Bw[i,j] = Bw[i,j] - Bw[i,k]*B_k[j]\n",
    "                end\n",
    "                Bw[i,k] = 0\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    Bw\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18c688ac",
   "metadata": {},
   "source": [
    "### Cyclic partitioning: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba555b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "function ge_par_cyclic!(B,n,m,comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nw = div(n,nranks)\n",
    "    # Init buffers\n",
    "    T = eltype(B)\n",
    "    if rank == 0\n",
    "        buffer_root = Vector{T}(undef,n*m)\n",
    "        # Fill buffer cyclicly\n",
    "        for i in 1:n\n",
    "            ub = i * m \n",
    "            lb = ub - m + 1\n",
    "            j = mod(i-1,nw)*nranks + div(i-1,nw) +1\n",
    "            buffer_root[lb:ub] = transpose(B)[:,j]\n",
    "        end\n",
    "    else\n",
    "        buffer_root = Vector{T}(undef,0)\n",
    "    end    \n",
    "    buffer =  Vector{T}(undef,nw*m)\n",
    "    # Send local matrices to workers\n",
    "    MPI.Scatter!(buffer_root,buffer,comm;root=0)\n",
    "    Bw = Matrix{T}(undef,nw,m)\n",
    "    transpose(Bw)[:] = buffer\n",
    "    MPI.Barrier(comm)\n",
    "    # time calcultation\n",
    "     t = @elapsed ge_cyclic_worker!(Bw,n,comm)\n",
    "    # Gather results\n",
    "     buffer[:] = transpose(Bw)[:]\n",
    "     MPI.Gather!(buffer,buffer_root,comm;root=0)\n",
    "     Tf = typeof(t)\n",
    "     if rank == 0\n",
    "         for i in 1:n\n",
    "             ub = i * m \n",
    "             lb = ub - m + 1\n",
    "             j = mod(i-1,nw)*nranks + div(i-1,nw) +1\n",
    "             transpose(B)[:,j] = buffer_root[lb:ub]\n",
    "         end\n",
    "         ts = Vector{Tf}(undef,nranks)\n",
    "     else\n",
    "         ts = Vector{Tf}(undef,0)\n",
    "     end\n",
    "     MPI.Gather!([t],ts,comm;root=0)\n",
    "    B, ts\n",
    "end\n",
    "\n",
    "function ge_cyclic_worker!(Bw,n,comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    nw,m = size(Bw)\n",
    "    B_k = similar(Bw,m)\n",
    "    my_rows = [i*nranks+(rank+1) for i in 0:(nw-1)]\n",
    "    @inbounds for k in 1:n\n",
    "        # If I have row k\n",
    "        if k in my_rows\n",
    "            myk = findfirst(my_rows .== k)[2]\n",
    "            # Update just row k\n",
    "            for t in (k+1):m\n",
    "                Bw[myk,t] =  Bw[myk,t]/Bw[myk,k]\n",
    "            end\n",
    "            Bw[myk,k] = 1\n",
    "            # Send k to other procs\n",
    "            B_k .= view(Bw,myk,:)\n",
    "            MPI.Bcast!(B_k,comm;root=rank)\n",
    "        else\n",
    "            owner = mod(k-1,nranks)\n",
    "            MPI.Bcast!(B_k,comm;root=owner)\n",
    "        end\n",
    "        # Do nothing if I only have rows < k \n",
    "        if k < maximum(my_rows)\n",
    "            cutoff = findfirst(my_rows .> k)[2]\n",
    "            for i in cutoff:length(my_rows)\n",
    "                for j in (k+1):m\n",
    "                    Bw[i,j] = Bw[i,j] - Bw[i,k]*B_k[j]\n",
    "                end\n",
    "                Bw[i,k] = 0\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    Bw\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe0bb4ab",
   "metadata": {},
   "source": [
    "### Main Function to run code and test performance: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d00c192",
   "metadata": {},
   "outputs": [],
   "source": [
    "using MPI\n",
    "using Test\n",
    "MPI.Init()\n",
    "\n",
    "\n",
    "function tridiagonal_matrix(n)\n",
    "    C = zeros(n,n)\n",
    "    stencil = [(-1,2,-1),(-1,0,1)]\n",
    "    for i in 1:n\n",
    "        for  (coeff,o) in zip((-1,2,-1),(-1,0,1))\n",
    "            j = i+o\n",
    "            if j in 1:n\n",
    "                C[i,j] = coeff\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    C\n",
    "end\n",
    "\n",
    "function ge_seq!(B)\n",
    "    n,m = size(B)\n",
    "    @inbounds for k in 1:n\n",
    "        for t in (k+1):m\n",
    "            B[k,t] =  B[k,t]/B[k,k]\n",
    "        end\n",
    "        B[k,k] = 1\n",
    "        for i in (k+1):n \n",
    "            for j in (k+1):m\n",
    "                B[i,j] = B[i,j] - B[i,k]*B[k,j]\n",
    "            end\n",
    "            B[i,k] = 0\n",
    "        end\n",
    "    end\n",
    "    B\n",
    "end\n",
    "\n",
    "function main_check()\n",
    "    # Create communicator \n",
    "    comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    n = 10*nranks\n",
    "    # create matrix B such that mod(n, nprocs) == 0\n",
    "    if rank == 0 \n",
    "        A = tridiagonal_matrix(n)\n",
    "        b = ones(n)\n",
    "        B = [A b]\n",
    "    else \n",
    "        B = tridiagonal_matrix(0)\n",
    "    end\n",
    "    # call parallel method \n",
    "    B_block, ts_block = ge_par_block!(copy(B),n,n+1,comm)\n",
    "    B_cyclic, ts_cyclic = ge_par_cyclic!(copy(B),n,n+1,comm)\n",
    "    # test if results is equal to sequential \n",
    "    if rank == 0\n",
    "        B_seq = ge_seq!(copy(B))\n",
    "        @test B_seq == B_block\n",
    "        @test B_seq == B_cyclic\n",
    "        # print timings\n",
    "        println(\"No workers: $nranks| Blockwise: min $(minimum(ts_block))| Cyclic: min $(minimum(ts_cyclic))\")\n",
    "    end\n",
    "end\n",
    "\n",
    "main_check()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbbd3b8f",
   "metadata": {},
   "source": [
    "## ASP : Exercise 1\n",
    "Implementation of parallel ASP using only `MPI.Bcast!`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec5f932f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function floyd_worker_bcast!(Cw,comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    Nw,N = size(Cw)\n",
    "    C_k = similar(Cw,N)\n",
    "    lb = Nw*rank + 1\n",
    "    ub = Nw*(rank+1)\n",
    "    for k in 1:N\n",
    "        if k in lb:ub\n",
    "            myk = k - lb + 1\n",
    "            C_k .= view(Cw,myk,:)\n",
    "            MPI.Bcast!(C_k,comm;root=rank)\n",
    "        else\n",
    "            owner = div(k-1,Nw)\n",
    "            MPI.Bcast!(C_k,comm;root=owner)\n",
    "        end\n",
    "        for j in 1:N\n",
    "            for i in 1:Nw\n",
    "                @inbounds Cw[i,j] = min(Cw[i,j],Cw[i,k]+C_k[j])\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    Cw\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19641daf",
   "metadata": {},
   "source": [
    "### Exercise: Measure search overhead\n",
    "Modify the code of the serial and parallel algorithms so that the functions return the number of nodes in the search tree that they visit. You can then compare how many more nodes are visited by the parallel algorithm compared with the serial algorithm (known as _search overhead_). You can then use the third cell to gather some statistics about the search overhead using your altered version of the functions. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22ccce39",
   "metadata": {},
   "source": [
    "## TSP: Exercise x (Measure search overhead)\n",
    "This is the solution of how the code can be altered to measure the search overhead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4d5ab70",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Distributed\n",
    "\n",
    "if procs() == workers()\n",
    "    addprocs(4)\n",
    "end\n",
    "\n",
    "@everywhere function visited(city,hops,path)\n",
    "    for i = 1:hops\n",
    "        if path[i] == city\n",
    "            return true\n",
    "        end\n",
    "    end\n",
    "    return false\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcee99f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TSP serial \n",
    "function tsp_serial_impl(connections,hops,path,current_distance, min_path, min_distance, node_count)\n",
    "    num_cities = length(connections)\n",
    "    if hops == num_cities\n",
    "        if current_distance < min_distance\n",
    "            min_path .= path\n",
    "            return min_path, current_distance, node_count\n",
    "        end\n",
    "    else\n",
    "        current_city = path[hops]\n",
    "        next_hops = hops + 1\n",
    "        for (next_city,distance_increment) in connections[current_city]\n",
    "            if !visited(next_city,hops,path)\n",
    "                node_count += 1\n",
    "                path[next_hops] = next_city\n",
    "                next_distance = current_distance + distance_increment\n",
    "                if next_distance < min_distance\n",
    "                    min_path, min_distance, node_count = tsp_serial_impl(connections,next_hops,path,next_distance,min_path,min_distance, node_count)\n",
    "                end\n",
    "            end\n",
    "        end        \n",
    "    end\n",
    "    return min_path, min_distance, node_count\n",
    "end\n",
    "\n",
    "function tsp_serial(connections,city)\n",
    "    num_cities = length(connections)\n",
    "    path=zeros(Int,num_cities)\n",
    "    hops = 1\n",
    "    path[hops] = city\n",
    "    min_path = zeros(Int, num_cities)\n",
    "    current_distance = 0\n",
    "    min_distance = typemax(Int)\n",
    "    node_count = 1\n",
    "    # Count the number of nodes visited in recursive function and return\n",
    "    min_path, min_distance, node_count = tsp_serial_impl(connections,hops,path,current_distance, min_path, min_distance, node_count)\n",
    "    (;path=min_path,distance=min_distance, node_count)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "327f5349",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TSP distributed\n",
    "@everywhere function tsp_dist_impl(connections,hops,path,current_distance,min_dist_chnl, max_hops,jobs_chnl,ftr_result,node_count)\n",
    "    num_cities = length(connections)\n",
    "    if hops == num_cities\n",
    "        min_distance = fetch(min_dist_chnl)\n",
    "        if current_distance < min_distance\n",
    "            take!(min_dist_chnl)\n",
    "            if ftr_result !== nothing\n",
    "                @spawnat 1 begin\n",
    "                    result = fetch(ftr_result)\n",
    "                    result.path .= path\n",
    "                    result.min_distance_ref[] = current_distance\n",
    "                end |> wait\n",
    "            end\n",
    "            put!(min_dist_chnl, current_distance)\n",
    "        end\n",
    "    elseif hops <= max_hops\n",
    "        current_city = path[hops]\n",
    "        next_hops = hops + 1\n",
    "        for (next_city,distance_increment) in connections[current_city]\n",
    "            if !visited(next_city,hops,path)\n",
    "                path[next_hops] = next_city\n",
    "                node_count += 1\n",
    "                next_distance = current_distance + distance_increment\n",
    "                min_distance = fetch(min_dist_chnl)\n",
    "                if next_distance < min_distance\n",
    "                    node_count = tsp_dist_impl(connections,next_hops,path,next_distance,min_dist_chnl,max_hops,jobs_chnl,ftr_result,node_count)\n",
    "                end\n",
    "            end\n",
    "        end \n",
    "    else\n",
    "        if jobs_chnl !== nothing \n",
    "            path_copy = copy(path) \n",
    "            put!(jobs_chnl,(;hops,path=path_copy,current_distance))\n",
    "        end\n",
    "    end\n",
    "    return node_count\n",
    "end\n",
    "\n",
    "function tsp_dist(connections,city)\n",
    "    max_hops = 2\n",
    "    num_cities = length(connections)\n",
    "    path=zeros(Int,num_cities)\n",
    "    result_path=zeros(Int, num_cities)\n",
    "    hops = 1\n",
    "    path[hops] = city\n",
    "    current_distance = 0\n",
    "    min_distance = typemax(Int)\n",
    "    node_count = 1\n",
    "    jobs_chnl = RemoteChannel(()->Channel{Any}(10))\n",
    "    min_dist_chnl = RemoteChannel(()->Channel{Int}(1))\n",
    "    put!(min_dist_chnl, min_distance)\n",
    "    ftr_result = @spawnat 1 (;path=result_path,min_distance_ref=Ref(min_distance))\n",
    "    # Add another future to store number of visited nodes\n",
    "    ftr_node_count = @spawnat 1 node_count_ref = Ref(node_count)\n",
    "    @async begin\n",
    "        ncount = 0\n",
    "        ncount += tsp_dist_impl(connections,hops,path,current_distance,min_dist_chnl,max_hops,jobs_chnl,nothing, ncount)\n",
    "        # Update node counter\n",
    "        node_count_ref = fetch(ftr_node_count)\n",
    "        node_count_ref[] += ncount \n",
    "        for w in workers()\n",
    "            put!(jobs_chnl,nothing)\n",
    "        end\n",
    "    end\n",
    "    @sync for w in workers()\n",
    "        @spawnat w begin\n",
    "            path = zeros(Int, num_cities)\n",
    "            max_hops = typemax(Int)\n",
    "            while true\n",
    "                job = take!(jobs_chnl)\n",
    "                if job == nothing\n",
    "                    break\n",
    "                end\n",
    "                hops = job.hops\n",
    "                path = job.path \n",
    "                current_distance = job.current_distance\n",
    "                ncount = 0\n",
    "                min_distance = fetch(min_dist_chnl)\n",
    "                if current_distance < min_distance\n",
    "                    ncount += tsp_dist_impl(connections,hops,path,current_distance,min_dist_chnl,max_hops,nothing,ftr_result, ncount)\n",
    "                    # Update node counter\n",
    "                    @spawnat 1 begin \n",
    "                        node_count_ref = fetch(ftr_node_count)\n",
    "                        node_count_ref[] += ncount \n",
    "                    end\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end \n",
    "    result = fetch(ftr_result)\n",
    "    # Fetch number of visited nodes for return\n",
    "    node_count_ref = fetch(ftr_node_count)\n",
    "    (;path = result.path, distance = result.min_distance_ref[], node_count=node_count_ref[])\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "706242f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Distributed\n",
    "using Plots\n",
    "using Statistics\n",
    "\n",
    "# Generate random matrices\n",
    "function rand_symmetric_distance_table(n)\n",
    "    threshold = 0.2\n",
    "    mincost = 3\n",
    "    maxcost = 10\n",
    "    infinity = 10000*maxcost\n",
    "    C = fill(infinity,n,n)\n",
    "    for j in 1:n\n",
    "      for i in 1:j\n",
    "        if rand() > threshold\n",
    "          C[i,j] = rand(mincost:maxcost)\n",
    "          C[j,i] = C[i,j]\n",
    "        end\n",
    "      end\n",
    "      C[j,j] = 0\n",
    "    end\n",
    "    C,infinity\n",
    "end\n",
    "\n",
    "# Sort to get ascending distance weights\n",
    "function sort_neighbors(C)\n",
    "    n = size(C,1)\n",
    "    map(1:n) do i\n",
    "        Ci = C[i,:]\n",
    "        cities = sortperm(Ci)\n",
    "        distances = Ci[cities]\n",
    "        collect(zip(cities,distances))[2:end]\n",
    "    end\n",
    "end\n",
    "\n",
    "n_cities = [4, 6, 8, 10]\n",
    "n_rep = 10\n",
    "city = 1\n",
    "node_count_serial = zeros(Union{Missing,Int64},length(n_cities), n_rep)\n",
    "node_count_dist = zeros(Union{Missing,Int64},length(n_cities), n_rep)\n",
    "for (i, n) in enumerate(n_cities)\n",
    "    @show n\n",
    "    for j in 1:n_rep\n",
    "        # Generate random connections matrix\n",
    "        C, inf = rand_symmetric_distance_table(n)\n",
    "        C = sort_neighbors(C)\n",
    "        # Run serial algorithm\n",
    "        path, distance, ncount_serial = tsp_serial(C, city)\n",
    "        # Check if graph is connected \n",
    "        if  distance >= inf\n",
    "            println(\"The input graph size $n, it $j is not connected\")\n",
    "            node_count_serial[i,j] = missing\n",
    "            node_count_dist[i,j] = missing\n",
    "        else\n",
    "            path, distance, ncount_dist = tsp_dist(C, city)\n",
    "            node_count_serial[i,j] = ncount_serial\n",
    "            node_count_dist[i,j] = ncount_dist\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "# Calculate average and confidence interval\n",
    "search_overhead_perc = (node_count_dist .- node_count_serial)./node_count_serial\n",
    "avg_search_overhead = [mean(skipmissing(search_overhead_perc[i,:])) for i in axes(search_overhead_perc,1)]\n",
    "conf_int = [1.96*std(skipmissing(search_overhead_perc[i,:]))/\n",
    "        sqrt(count(!ismissing,search_overhead_perc[i,:])) \n",
    "        for i in axes(search_overhead_perc,1)]\n",
    "\n",
    "# Plot\n",
    "plot(n_cities, avg_search_overhead, ribbon=conf_int, markershape=:circle, legend=false)\n",
    "title!(\"Average search overhead (%) and 95%-CI \\nin parallel algorithm\")\n",
    "ylabel!(\"Average % extra nodes visited \\n in parallel algorithm\")\n",
    "xlabel!(\"Number of cities\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "968304a6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.1",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
