// Jun 27 2017 13:12:10
// Copyright 2017 Brendan Halpin

// pam: Partitioning Around Medoids

// Newvar: variable holding the cluster solution
// Mandatory Options:
// - DISTMAT: pairwise distance matrix
// - IDVAR: ID variable to verify sort order of data and DISTMAT match
// - MEDOIDS: Either the number of clusters to fit, or a binary variable indicating which cases are medoids
// Optional options:
// - MANY or GA: mutually exclusive, if both are absent start depends on medoids() (if number, random)
//   - MANY: Start from multiple random starting points, select best
//   - GA: use a genetic algorithm to find an approximate best solution (GAPOP changes the population size)
// - NONVORONOI: use slightly slower non-Voronoi algorithm

program clpam, rclass
syntax newvarlist(max=1), DISTmat(string) IDvar(varname) MEDoids(string) [ MANY GA NONVORonoi GAPop(real 100)]

// Determine whether medoids option is a variable or a number
local medoidstype "undetermined"
capture su `medoids'
if (_rc == 0) {
  qui su `medoids'
  di "Medoids (n=`r(sum)') taken from `medoids'"
  local medoidstype "var"
}
else {
  if (_rc==111) {
    di "medoids var doesn't exist"
  }
  else {
    if (real("`medoids'")+0==`medoids') {
      di "Random starting medoids (Nk=`=floor(`medoids')')"
      local medoidstype "num"
    }
    else {
      di "Something is wrong with medoid variable"
    }
  }
}

local nonv 0
if ("`nonvoronoi'"!="") {
  local nonv 1
}

// Check dist-mat exists
qui matlist `distmat'[1,1]

// Insist on correct sort order, and that it is unique
isid `idvar', sort

// Initialise variable relating to memoisation in the top-level environment
mata: ga_memo = .
mata: ga_hit = .
mata: ga_insert = .

mata: pam_temp_distmat = st_matrix("`distmat'")

if (("`many'"!="") & ("`ga'"!="")) {
  di as error "Only one of MANY and GA may be selected"
  exit
}
local searchtype 1
if ("`many'"!="") {
  local searchtype 2
}
if ("`ga'"!="") {
  local searchtype 3
}

if (("`searchtype'"=="1") | ("`searchtype'"=="")) {
  if ("`medoidstype'"=="var") {
    mata: pam_temp_medoids = pam_vanilla(pam_temp_distmat, medoids_from_var("`medoids'"),`nonv')
  }
  else {
    mata: pam_temp_medoids = pam_vanilla(pam_temp_distmat, random_medoids(`medoids', pam_temp_distmat), `nonv')
  }
}
else if ("`searchtype'" == "2") {
  di "Trying multiple starting points"
  if ("`medoidstype'"=="var") {
    di as error "PAM with multiple starting points (option MANY) is not compatible with specifying starting medoids"
    exit
  }
  else {
    mata: pam_temp_medoids = pam_vanilla(pam_temp_distmat, random_medoids_many(`medoids', pam_temp_distmat), `nonv')
  }
}
else if ("`searchtype'"=="3") {
  if ("`medoidstype'"=="var") {
    di as error "PAM with option GA is not compatible with specifying starting medoids"
    exit
  }
  else {
    mata: pam_temp_medoids = pam_ga(pam_temp_distmat, `medoids', `gapop', `nonv')[1,]
  }
}

// Return the medoids
tempname ret_med
mata: st_matrix("`ret_med'",pam_temp_medoids)
return matrix medoids = `ret_med'

mata: pam_temp_group = getgroup(pam_temp_distmat, pam_temp_medoids)

getmata `varlist'=pam_temp_group

mata: mata drop pam_temp_distmat pam_temp_medoids pam_temp_group ga_memo ga_hit ga_insert

end // program pam

capture mata mata drop getgroup()
mata
real matrix getgroup (real matrix distmat, real matrix medoids) {
  nearestvalue = rowmin(distmat[.,medoids])
  nearestcase  = distmat[.,medoids] :== nearestvalue
  group = rowmax(nearestcase:*transposeonly(range(1,length(medoids),1)))
  return(group)
}
end

capture mata mata drop medoids_from_var()
mata
real matrix medoids_from_var(string varname) {
  V = .
  st_view(V,.,varname)
  return(select(range(1,length(V),1), V))
}
end

capture mata mata drop reassign_medoids()
mata
real matrix function reassign_medoids (matrix medoids, matrix distmat) {
  // Given a set of medoids, identify sets of elements closest to each
  // medoid, then return the set of true medoids of these groups
  N = rows(distmat)
  Nk = length(medoids)
  group = getgroup(distmat,medoids)
  // For each group find the true medoid
  newmedoids = J(Nk,1,.)
  for (i=1;i<=Nk;i++) {
    // Select the cluster
    selector = select(range(1,N,1), (group :== i))
    clusterdistmat = distmat[selector, selector]

    clustersize = rows(clusterdistmat)
    centroiddist = mean(clusterdistmat)
    mindist = min(centroiddist)

    for (row = 1; row <= clustersize; row++) {
      if (mindist == centroiddist[row]) {
        newmedoids[i] = selector[row] // reverse select
      }
    }
  }
  newmedoids = uniqrows(newmedoids)
  return(newmedoids[order(newmedoids,1)])
}
end

capture mata mata drop reassign_medoids_non_voronoi()
mata
real matrix function reassign_medoids_non_voronoi (matrix medoids, matrix distmat) {
  // Given a set of medoids, identify sets of elements closest to each
  // medoid, then return the set of revised medoids of these groups
  // Non-voronoi: take first random medoid that improves
  N = rows(distmat)
  Nk = length(medoids)
  group = getgroup(distmat,medoids)
  // For each group find an improved medoid
  newmedoids = medoids //J(Nk,1,.)
  for (i=1;i<=Nk;i++) {
    // Select the cluster rows, but all columns
    selector = select(range(1,N,1), (group :== i))
    clusterdistmat = distmat[selector,]

    centroiddist = mean(clusterdistmat)
    maxdist = max(centroiddist)
    newdist = maxdist+1;

    for (j=1; j<=min((N,100)); j++) { // note break below
      newmedoid = runiformint(1,1,1,N) // 1  + floor(runiform(1,1)*N)
      newdist = centroiddist[newmedoid]
      if (newdist<centroiddist[medoids[i]]) {
        newmedoids[i] = newmedoid
        break
      }
    }
  }
  if ((medoids != newmedoids) & (min(distmat[newmedoids,newmedoids] + I(length(newmedoids))) != 0)) {
    medoids = newmedoids
  }
  return(medoids[order(medoids,1)])
}
end

capture mata mata drop random_medoids()
mata
real matrix function random_medoids(real Nk, real matrix distmat) {
  N = rows(distmat)
  // Initialise with random medoids (Nk)
  medoids = uniqrows(floor(runiform(Nk,1):*N) :+ 1)
  // Replace initial medoids if there are duplicates (zero on off-diag)
  while ( (length(medoids)<Nk) | (min(distmat[medoids,medoids] + I(length(medoids))) == 0)) {
    medoids = uniqrows(floor(runiform(Nk,1):*N) :+ 1)
  }
  return(medoids)
}
end

capture mata mata drop random_medoids_many()
mata
real matrix random_medoids_many(real Nk, real matrix distmat) {
  many = 100
  meds = J(many,Nk,.)
  for (i=1;i<=many; i++) {
    meds[i,] = transposeonly(random_medoids(Nk,distmat))
  }
  return(meds)
}
end

capture mata mata drop pam_sum_within_clusters()
mata
real scalar function pam_sum_within_clusters (real matrix distmat, real matrix medoids) {
  external ga_memo, ga_hit, ga_insert
  N = rows(distmat)
  Nk = length(medoids)
  group = getgroup(distmat,medoids)

  if (asarray_contains(ga_memo, transposeonly(medoids))) {
    SS = asarray(ga_memo, transposeonly(medoids))
    ga_hit++
  } else {
    // For each group find the true medoid
    SS = 0
    for (i=1;i<=Nk;i++) {
      // Select the cluster
      selector = select(range(1,N,1), (group :== i))
      clusterdistmat = distmat[selector, selector]
      SStemp =sum(clusterdistmat)/rows(clusterdistmat)
      SS = SS + SStemp
    }
    asarray(ga_memo, transposeonly(medoids), SS)
    ga_insert++
  }
  return(SS)
}

end

capture mata mata drop pam_iterate()
mata
real matrix pam_iterate (real matrix distmat, real matrix medoids, real scalar nonvor) {
  // With a proposed set of medoids, keep reassigning them until stability
  // Effectively a hill-climb function
  Nk=rows(medoids)
  oldmed = J(Nk,1,0)
  i=0
  while ((oldmed != medoids)) {
    i++
    oldmed = medoids
    // If defective, randomise
    if ( (length(medoids)<Nk) | (min(distmat[medoids,medoids] + I(length(medoids))) == 0)) {
      medoids = random_medoids(Nk, distmat)
    }
    if (nonvor) {
      medoids = reassign_medoids_non_voronoi(medoids, distmat)
    }
    else {
      medoids = reassign_medoids(medoids, distmat)
    }
  }
  return(medoids)
}
end

capture mata mata drop pam_ga()
mata
real matrix function pam_ga (real matrix distmat, real scalar Nk, real scalar npop, real scalar nonv) {
  external ga_memo, ga_hit, ga_insert

  // Memoise the fitness function
  // Is worth about 10-12% in time.
  // Used by pam_sum_within_clusters()
  ga_memo = asarray_create("real",Nk)
  ga_hit = 0
  ga_insert = 0

  npop = 2*floor(npop/2) // even
  nparents = floor(npop/2) // change from 50% to intensify selection
  nsurv = nparents
  nnew = floor(npop*0.20)
  ntoconv = ceil(npop*0.1)
  // Use top 50% as core. They procreate to replace 30%, and 20% is new random
  // Progress data is on top 50%

  // Create a random population, each row a medoid set
  population = J(npop,Nk,.)
  sigmadist = J(npop,1,.)
  for (i=1; i<=npop; i++) {
    population[i,.] = transposeonly(random_medoids(Nk, distmat))
//    sigmadist[i] = pam_sum_within_clusters(distmat, transposeonly(population[i,.]))
  }

  // Main iteration
  iter = 0
  while (iter==0 |(max(sigmadist[1..ntoconv]) != min(sigmadist[1..ntoconv]))) {
    newpop = J(npop,Nk,.)
    iter++
    // Replace each random medoid set with a PAM hill-climbed version
    for (i=1; i<=npop; i++) {
      medoids = transposeonly(population[i,.])
      newpop[i,.] = transposeonly(pam_iterate(distmat,medoids,nonv))
    }

    // Calculate fitness (lower is better)
    for (i=1; i<=npop; i++) {
      medoids = transposeonly(newpop[i,.])
      sigmadist[i] = pam_sum_within_clusters(distmat, medoids)
    }

    // Order the new medoid sets per fitness (best (lowest) first)
    newpop = newpop[order(sigmadist,1),][1..nparents,]
    sigmadist = sigmadist[order(sigmadist,1)]
    "Iter "+strofreal(iter,"%-3.0f")+": Σ dist: min "+strofreal(min(sigmadist[1..ntoconv]))+"; mean: "+strofreal(mean(sigmadist[1..ntoconv]))+"; max: "+strofreal(max(sigmadist[1..ntoconv]))
    displayflush()

    // Keep top 50%
    population[1..nsurv,] = newpop[1..nsurv,]

    // Crossover section
    // Top 50% used for crossover to create 30%, remaining 20% new
    for (i=1+nsurv; i<=npop-nnew; i++) {
      j = 1+floor((1-runiform(1,1)^0.1)*nparents)
      k = 1+floor((1-runiform(1,1)^0.1)*nparents)
      l = runiformint(1,1,1,Nk-1)
      population[i,1..l] = newpop[j,1..l]
      population[i,l+1..Nk] = newpop[k,l+1..Nk]
    }
    for (i=1+npop-nnew; i<=npop; i++) {
      if (runiform(1,1)<0.1) {
        population[i,.] = transposeonly(random_medoids(Nk, distmat))
      }
    }
  }
  sprintf("Memoisation: hits: %1.0f; inserts: %1.0f.", ga_hit,ga_insert)
  return(population)
}
end

capture mata mata drop pam_vanilla()
mata
real matrix function pam_vanilla(real matrix distmat, real medoids, real scalar nonv) {
  external ga_memo, ga_hit, ga_insert

  // If medoids is a matrix, each row is a medoid set
  // If it's a vector, it's a col
  many_meds = min((rows(medoids),cols(medoids)))
  if (many_meds>1) {
//    "many"
    Nk = length(transposeonly(medoids[1,]))
  }
  else {
//    "few"
    Nk = length(medoids)
  }

  // Memoise the fitness function
  // Is worth about 10-12% in time.
  // Used by pam_sum_within_clusters()
  ga_memo = asarray_create("real",Nk)
  ga_hit = 0
  ga_insert = 0

  if (many_meds>1) {
    results = J(many_meds,length(medoids[1,]),.)
    scores  = J(many_meds,1,.)
    for (i=1; i<=many_meds; i++) {
      results[i,1..Nk] = transposeonly(pam_iterate(distmat,transposeonly(medoids[i,]),nonv))
      scores[i] = pam_sum_within_clusters(distmat, transposeonly(medoids[i,]))
    }
    result=results[order(scores,1),][1,]
  }
  else {
    result = pam_iterate(distmat,medoids,nonv)
  }
  return(result)
}
end