# Force directed graphs, working first on the horizontal axis,
# then the vertical.

library(dplyr)
library(readr)
library(ggplot2)
library(ggrepel)

plot_intermediate_conv = FALSE
plot_intermediate = FALSE
num_trials = 50

# Choices are "max", "avg", "max_asym", "max_asym_wtd":
blank_types = c("max", "avg", "max_asym", "max_asym_wtd")
distance_powers = c(1, 3)
states = c("nsw", "vic", "qld", "wa", "sa", "tas", "act", "nt")


png_print = function(out_file, img_width, img_height, ggplot_img) {
  # A function to print a plot to png:
  png(filename=out_file, width=img_width, height=img_height, units="px")
  print(ggplot_img)
  dev.off()
  return(TRUE)
}

make_plot = function(x, group_names, HTV, htv_colours, title_str, x_lim=NULL, y_lim=NULL) {
  locations.df = data.frame(x = x[, 1], y = x[, 2], names=group_names, HTV = HTV)
  
  scatter_plot = ggplot() + 
    geom_point(data=locations.df, aes(x=x, y=y, colour=factor(HTV))) +
    geom_text_repel(data=locations.df, aes(x=x, y=y, label=names, colour=factor(HTV)),
                    size=6, show.legend=FALSE) +
    scale_x_continuous(limits=x_lim) + scale_y_continuous(limits=y_lim) +
    coord_fixed(ratio = 1) +
    scale_color_manual(values=htv_colours, name="HTV cards") +
    theme_grey(base_size = 16) +
    guides(colour = guide_legend(override.aes = list(size=4))) +
    theme(legend.text=element_text(size=16)) +
    ggtitle(title_str)
  
  return(scatter_plot)
}


for (distance_power in distance_powers) {
  for (blank_type in blank_types) {
    
    output_folder = sprintf("images/force_directed/trials/%s_%d_seq", blank_type, distance_power)
    dir.create(output_folder, showWarnings = FALSE)
    
    for (state in states) {
      print(sprintf("***** Starting %s %s %d *****", state, blank_type, distance_power))
      
      source("htv_abbrevs.R")
      htv_colours = c(ALP = "#FF0000", Lib = "#0000FF", Both = "#800080", Neither = "#000000")
      
      if (state %in% c("act", "nt")) {
        alp_htv = alp_htv[1:3]
        lib_htv = lib_htv[1:3]
      }
      
      in_file = sprintf("graph_distances/distances_max_%s_blank%s.csv", state, blank_type)
      
      distances.df = read_csv(in_file)
      
      group_names = names(distances.df)
      num_groups = length(group_names)
      num_dims = 2
      
      ALP_i = which(names(distances.df) == "ALP")
      LNP_i = which(names(distances.df) == lib_abbrev)
      PHON_i = which(names(distances.df) == phon_rua)
      
      ALP_HTV = character(num_groups)
      Lib_HTV = character(num_groups)
      
      ALP_HTV[which(group_names %in% alp_htv)] = "ALP"
      Lib_HTV[which(group_names %in% lib_htv)] = "Lib"
      
      HTV = sprintf("%s%s", ALP_HTV, Lib_HTV)
      HTV = gsub("ALPLib", "Both", HTV)
      HTV[which(HTV == "")] = "Neither"
      
      title_str = sprintf("Force-directed, %s,\nsequential, %s, d^%d", state_title, blank_type, distance_power)
      
      
      dist_matrix = as.matrix(distances.df)
      dist_matrix = dist_matrix / max(dist_matrix)
      dist_matrix = dist_matrix^distance_power
      
      x = matrix(0, nrow=num_groups, ncol = num_dims)
      mean_x = 0*x
      
      attraction_const = 1
      inc = 5e-2
      tol = 1e-6
      
      max_iter = 1000000
      list_x = list()
      
      # First row min, second row max:
      axis_bounds = matrix(0, nrow=2, ncol=num_dims)
      
      for (i_trial in 1:num_trials) {
        print(sprintf("Starting trial %d", i_trial))
        
        for (i_dim in 1:num_dims) {
          x[ , i_dim] = runif(num_groups)
          
          for (i_iter in 1:max_iter) {
            if (i_iter %% 10000 == 0) {
              print(i_iter)
            }
            
            prev_x = x
            
            r = list()
            length(r) = num_dims
            
            for (i in 1:num_dims) {
              r[[i]] = outer(x[ , i], rep(1, num_groups)) - outer(rep(1, num_groups), x[ , i])
              
              if (i == 1) {
                current_dist = r[[i]]*r[[i]]
              } else {
                current_dist = current_dist + r[[i]]*r[[i]]
              }
            }
            
            current_dist = sqrt(current_dist)
            
            spring_signs = 2*(current_dist > dist_matrix) - 1
            
            
            F_springs_i = -attraction_const * spring_signs * (abs(current_dist - dist_matrix)) * r[[i_dim]]
            F_i = rowSums(F_springs_i)
            x[ , i_dim] = x[ , i_dim] + inc*F_i
            
            if (plot_intermediate_conv) {
              if (i_iter %% 1000 == 0) {
                out_file = sprintf("images/force_directed/convergence/%s_seq_%s_%07d.png",
                                   state, blank_type, i_iter)
                
                scatter_plot = make_plot(x, group_names, HTV, htv_colours, title_str)
                png_print(out_file, 550, 500, scatter_plot)
              }
            }
            
            if (abs(max(x - prev_x)) < tol) {
              break
            }
          }
        }
        
        
        # Translate and rotate.
        
        ALP_pos = as.vector(x[ALP_i, ])
        LNP_pos = as.vector(x[LNP_i, ])
        
        mid_x = 0.5*(ALP_pos[1] + LNP_pos[1])
        mid_y = 0.5*(ALP_pos[2] + LNP_pos[2])
        
        delta_x = LNP_pos[1] - ALP_pos[1]
        delta_y = LNP_pos[2] - ALP_pos[2]
        
        x[ , 1] = x[ , 1] - mid_x
        x[ , 2] = x[ , 2] - mid_y
        
        theta = atan2(delta_y, delta_x)
        
        rot_matrix = matrix(0, nrow=2, ncol=2)
        rot_matrix[1, ] = c(cos(theta), sin(theta))
        rot_matrix[2, ] = c(-sin(theta), cos(theta))
        
        x = t(rot_matrix %*% t(x))
        
        if (x[PHON_i, 2] < 0) {
          x[ , 2] = -x[ , 2]
        }
        
        
        if (i_trial == 1) {
          mean_x = x
        } else {
          # The PHON_i lines above are supposed to ensure that
          # One Nation appears in the upper half-plane, but maybe
          # it's close to y=0 and it ended up in the lower half-plane.
          # So, check if the current arrangement is close to the current one.
          
          test_x = x
          test_x[ , 2] = -test_x[ , 2]
          
          diff_1 = sum(abs(mean_x - x))
          diff_2 = sum(abs(mean_x - test_x))
          
          if (diff_2 < diff_1) {
            x = test_x
          }
          
          mean_x = mean_x + x
        }
        
        
        out_file_png = sprintf("%s/%s_%04d.png", output_folder, state, i_trial)
        out_file_csv = sprintf("%s/%s_%04d.csv", output_folder, state, i_trial)
        
        if (plot_intermediate) {
          scatter_plot = make_plot(x, group_names, HTV, htv_colours, title_str)
          png_print(out_file_png, 600, 500, scatter_plot)
        }
        
        list_x[[i_trial]] = x
        
        if (i_trial == 1) {
          for (i in 1:num_dims) {
            axis_bounds[1, i] = min(x[ , i])
            axis_bounds[2, i] = max(x[ , i])
          }
        } else {
          for (i in 1:num_dims) {
            this_min = min(x[, i])
            this_max = max(x[ ,i])
            
            if (this_min < axis_bounds[1, i]) {
              axis_bounds[1, i] = this_min
            }
            
            if (this_max > axis_bounds[2, i]) {
              axis_bounds[2, i] = this_max
            }
          }
        }
        
        write(t(x), file=out_file_csv, ncolumns= num_dims, sep=",")
      }
      
      # Re-print the intermediate plots, now with fixed axis extents.
      x_bounds = c(axis_bounds[1, 1], axis_bounds[2, 1])
      y_bounds = c(axis_bounds[1, 2], axis_bounds[2, 2])
      
      for (i_trial in 1:num_trials) {
        x = list_x[[i_trial]]
        out_file_png = sprintf("%s/%s_%04d.png", output_folder, state, i_trial)
        scatter_plot = make_plot(x, group_names, HTV, htv_colours, title_str, x_bounds, y_bounds)
        png_print(out_file_png, 600, 500, scatter_plot)
      }
      
      
      mean_x = mean_x / num_trials
      
      out_file = sprintf("images/force_directed/%s_%s_%d_seq.png", state, blank_type, distance_power)
      
      final_plot = make_plot(mean_x, group_names, HTV, htv_colours, title_str)
      print(final_plot)
      png_print(out_file, 600, 500, final_plot)
    }
  }
}
