# 3D plot
#
# R. Perry, July 2018
#
# based on graphData() from CORVIDS/RecreateData.py

from numpy import ones
from random import sample
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # so projection='3d' will work in add_subplot()

def plot(all_sols, max_samples=20, vals=None):

  fig = plt.figure()
  ax1 = fig.add_subplot(111, projection='3d')

  if max_samples == -1 or len(all_sols) <= max_samples:
    sols = all_sols
  else:
    sols = sample(all_sols, min(max_samples, len(all_sols)))

  if len(sols) == 1:
    sols.append([0]*len(sols[0]))

  if vals == None:
    vals = [i+1 for i in range(len(sols[0]))]

  xpos = []
  ypos = []
  dz = []
  for index1, sol in enumerate(sols):
    for index2, val in enumerate(sol):
      ypos.append(vals[index2] - .5)
      xpos.append(index1 + .5)
      dz.append(val)

  num_elements = len(xpos)
  zpos = [0]*num_elements
  dx = ones(num_elements)
  dy = ones(num_elements)

  plt.ylim(min(ypos), max(ypos) + 1)

  ax1.bar3d(xpos, ypos, zpos, dx, dy, dz, color='#00ceaa')
  ax1.set_xlabel('Solution Number')
  ax1.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(base=1))
  ax1.set_ylabel('Response Value')
  ax1.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(base=1))
  ax1.set_zlabel('Frequency')
  ax1.zaxis.set_major_locator(matplotlib.ticker.MultipleLocator(base=1))
  plt.show()