module graphs

  use :: iso_c_binding, only : c_int, c_int64_t, c_double
  use :: fgsl
  use :: utilities
  use :: prob
  use :: sort

  implicit none
  
  private
  integer, parameter :: dp = c_double
  integer, parameter :: i4 = c_int
  integer, parameter :: i8 = c_int64_t

  type, public :: ContactGraph
    integer(kind=c_int64_t) :: order ! The number of nodes
    integer(kind=c_int64_t) :: size ! The number of edges
    integer(kind=c_int64_t), dimension(:), allocatable :: nodes
    integer(kind=c_int64_t), dimension(:), allocatable :: weights
    integer(kind=c_int64_t), dimension(:), allocatable :: degrees
    integer(kind=c_int64_t), dimension(:,:), allocatable :: edges
    integer(kind=c_int64_t) :: num_self_loops
    integer(kind=c_int64_t) :: rewiring_iterations
    integer(kind=c_int64_t), dimension(:,:), allocatable :: adjacency_matrix
  contains
    private
    procedure, pass(self) :: ChungLuGraph
    procedure, pass(self) :: FastChungLuGraph
    procedure, pass(self) :: ConfigurationGraph
    generic, public :: build => ConfigurationGraph
    procedure, pass(self) :: EdgesToAdjacency
    generic, public :: adjacency => EdgesToAdjacency
  end type ContactGraph
  
  interface ContactGraph
    module procedure init_ContactGraph
  end interface ContactGraph

contains

  function init_ContactGraph(nodes, weights) result(self)
    implicit none
    type(ContactGraph) :: self
    integer(kind=c_int64_t), dimension(:), intent(in) :: nodes
    integer(kind=c_int64_t), dimension(:), intent(in) :: weights

    if (size(nodes) .ne. size(weights)) then
      call FatalErrorMessage('in init_ContactGraph, size of nodes and weights must be the same')
    end if
  
    allocate(self%weights, source=weights)
    ! self%weights = weights
    allocate(self%nodes, source=nodes)

    self%order = size(nodes)

    ! self%nodes = nodes
    ! call Qsort(self%weights, self%nodes, -1)
  
    allocate(self%degrees, mold=nodes)
    self%degrees(:) = 0
  
    call self%build()

    call self%adjacency()
  
    return
  end function init_ContactGraph

  subroutine FastChungLuGraph(self)
    class(ContactGraph), intent(inout) :: self
  
    integer(kind=c_int64_t) :: N
    real(kind=c_double) :: S
  
    integer(kind=c_int64_t) :: u, v
    real(kind=c_double) :: p, q, r
  
    N = size(self%nodes)
  
    S = sum(self%weights)*1.0_dp
    do u=1,N-1
      v = u+1
      p = min(self%weights(u)*self%weights(v)/S, 1.0_dp)
      do while ((v<N+1) .and. (p>0))
        if ( abs(p - 1.0_dp) > meps ) then
          call pdists%unif_rvs(r, meps, 1.0_dp-meps)
          v = v + int(floor(log(r)/log(1.0_dp-p)))
        end if
        if (v<N+1) then
          q = min(self%weights(u)*self%weights(v)/S, 1.0_dp)
          call pdists%unif_rvs(r,meps,1.0_dp-meps)
          if (r<q/p) then
            if (allocated(self%edges)) then
              call push_row(self%edges,(/u,v/))
            else
              allocate(self%edges(1,2))
              self%edges(1,1) = u
              self%edges(1,2) = v
            end if
            self%degrees(u) = self%degrees(u) + 1
            self%degrees(v) = self%degrees(v) + 1
          end if
          p = q
          v = v+1
        end if
      end do
    end do

    self%size = size(self%edges(:,1))
  
    return
  end subroutine FastChungLuGraph

  subroutine ChungLuGraph(self)
    class(ContactGraph), intent(inout) :: self
  
    integer(kind=c_int64_t) :: N
    real(kind=c_double) :: S
    real(kind=c_double) :: c
  
    integer(kind=c_int64_t) :: u, v
    real(kind=c_double) :: p_uv, r
  
    N = size(self%nodes)
  
    S = sum(self%weights)*1.0_dp
    do u=1,N
      c = S - self%weights(u)
      do v=1,N
        p_uv = self%weights(u)*self%weights(v)/c
        if (u .ne. v) then
          call pdists%unif_rvs(r)
          if (r .le. p_uv) then
            ! Add edge to the graph
            if (allocated(self%edges)) then
              call push_row(self%edges,(/u,v/))
            else
              allocate(self%edges(1,2))
              self%edges(1,1) = u
              self%edges(1,2) = v
            end if
            self%degrees(u) = self%degrees(u) + 1
            self%degrees(v) = self%degrees(v) + 1
          end if
        end if
      end do
    end do

    self%size = size(self%edges(:,1))
  
    return
  end subroutine ChungLuGraph

  subroutine ConfigurationGraph(self)
    implicit none
    class(ContactGraph), intent(inout) :: self
  
    integer(kind=c_int64_t), dimension(:), allocatable :: edges_list, edges_list_shuffle
  
    integer(kind=c_int64_t), dimension(:), allocatable :: tmp
  
    integer :: j, k
    integer :: indxA, indxB
  
    integer(kind=c_int64_t) :: indx, u, v, tmp1, tmp2

    real(kind=c_double) :: p

    integer :: num_self_loops
    
    ! Check if the sum of the degrees is even
    if (mod(sum(self%weights),2)==1) then
      ! It isn't, so we have to change something.
      ! Why not add one to the degree of randomly selected node?
      ! Select a node at random:
      indx = pdists%randint(1,size(self%nodes))
      ! Add one to the degree of this node
      self%weights(indx) = self%weights(indx) + 1
    end if
  
    ! Make a list of all the half-stubs
    do j=1,size(self%nodes)
      allocate(tmp(self%weights(j)))
      tmp(:) = self%nodes(j)
      call push(edges_list, tmp)
      deallocate(tmp)
    end do
  
    ! Randomly shuffle the edges_list
    edges_list_shuffle = pdists%shuffle(edges_list)
  
    num_self_loops = 0
    ! Make edges from the pairs of the shuffled list
    do j=1,size(edges_list_shuffle)/2
      u = edges_list_shuffle(2*j-1)
      v = edges_list_shuffle(2*j)
      indxA = findloc(self%nodes, u, dim=1)
      indxB = findloc(self%nodes, v, dim=1)
      if (allocated(self%edges)) then
        call push_row(self%edges,(/u,v/))
      else
        allocate(self%edges(1,2))
        self%edges(1,1) = u
        self%edges(1,2) = v
      end if
      if (u==v) num_self_loops = num_self_loops + 1
      self%degrees(indxA) = self%degrees(indxA) + 1
      self%degrees(indxB) = self%degrees(indxB) + 1
    end do

    do k=1,100
      ! Loop through to find self_loops and rewire
      do j=1,size(self%edges(:,1))
        if (self%edges(j,1)==self%edges(j,2)) then
          ! This is a self-loop, so rewire
          indx = j
          ! Randomly choose an edge for swap
          do while (indx==j)
            indx = pdists%randint(1,size(self%edges(:,1)))
          end do
          ! Which rewiring to do?
          call pdists%unif_rvs(p)
          if (p<0.5_dp) then
            tmp1 = self%edges(j,1)
            tmp2 = self%edges(indx,1)
            self%edges(j,1) = tmp2
            self%edges(indx,1) = tmp1
          else
            tmp1 = self%edges(j,1)
            tmp2 = self%edges(indx,2)
            self%edges(j,1) = tmp2
            self%edges(indx,2) = tmp1
          end if
        end if
      end do

      num_self_loops = 0
      do j=1,size(self%edges(:,1))
        if (self%edges(j,1)==self%edges(j,2)) num_self_loops = num_self_loops + 1
      end do

      if (num_self_loops==0) exit
    end do

    self%size = size(self%edges(:,1))

    self%num_self_loops = num_self_loops
    self%rewiring_iterations = k
  
    return
  
  end subroutine ConfigurationGraph

  subroutine EdgesToAdjacency(self)
    implicit none
    class(ContactGraph), intent(inout) :: self

    integer(kind=c_int64_t) :: i, j, k
    integer(kind=c_int64_t) :: u, v

    if (.not. allocated(self%adjacency_matrix)) allocate(self%adjacency_matrix(self%order,self%order))
    self%adjacency_matrix(:,:) = 0

    do k=1,self%size
      u = self%edges(k,1)
      v = self%edges(k,2)

      i = findloc(self%nodes, u, dim=1)
      j = findloc(self%nodes, v, dim=1)

      self%adjacency_matrix(i,j) = self%adjacency_matrix(i,j) + 1_i8
      self%adjacency_matrix(j,i) = self%adjacency_matrix(j,i) + 1_i8
    end do

    return
  end subroutine EdgesToAdjacency
    


end module graphs