# Data

from scipy.spatial import Voronoi, voronoi_plot_2d, KDTree, distance
import matplotlib.pyplot as plt

class Pit:
  def __init__(self, loc, depth, diam):
    self.loc = loc
    self.depth = depth
    self.diam = diam

  def __repr__(self):
    return '%s @ %.1fcm deep, %.1fcm wide' % (str(self.loc), self.depth, self.diam)

  def disp(self):
    return '%.1fcm dp\n %.1fcm wd' % (self.depth, self.diam)

  def __getitem__(self,ind):
    return self.loc[ind]

class Date:
  def __init__(self, year, month, day):
    self.year = year
    self.month = month
    self.day = day
    
  def __repr__(self):
    return '%d-%d-%d' % (self.year, self.month, self.day)

class Trial:
  def __init__(self, date, intro, dead, size, pits):
    self.date = date
    self.intro = intro
    self.dead = dead
    self.size = size
    self.pits = pits
    self.pitlocs = [pit.loc for pit in self.pits]

  def __repr__(self):
    return str(self.date) + ' ' + str(self.pits)

  def plot(self, save=False):
    vor = Voronoi([pit.loc for pit in self.pits])
    voronoi_plot_2d(vor)
    for pit in self.pits:
      plt.text(pit[0], pit[1], pit.disp(), ha='center', va='bottom', size='small')
    plt.xlabel('%s (dimension %dx%dcm)' % (str(self.date), self.size[0], self.size[1]))
    if save:
      plt.savefig(str(self.date)+'.png', bbox_inches='tight')
    else:
      plt.show()

  def nearest_neighbor(self):
    tree = KDTree(self.pitlocs)
    sumnn = 0
    return [dists[1] for dists in tree.query(self.pitlocs,2)[0]]

trials = [
  Trial(Date(2019, 10, 16), 31, 6, [33,32], [
    Pit([4,25],1.3,4.2),
    Pit([3,13],1.4,3.7),
    Pit([10,25],1.1,3.0),
    Pit([18,18],1.8,2.3),
    Pit([30,14],2.2,3.1),
    Pit([29,11],1.4,2.5),
    Pit([27,10],1.2,2.1),
    Pit([29,6],2.4,3.9),
    Pit([26,5],1.8,3.6),
  ]),
  Trial(Date(2019, 10, 30), 27, 3, [24,24], [
    Pit([4,21],2.0,7.0),
    Pit([7,22],2.5,4.1),
    Pit([2,9],0.5,2.0),
    Pit([12,18],1.2,2.5),
    Pit([12,11],1.2,3.0),
    Pit([20,11],1.0,3.0),
    Pit([19,2],1.5,4.0),
  ]),
  Trial(Date(2019, 12, 3), 19, 3, [17, 16], [
    Pit([14,5],1.3,4.1),
    Pit([12,2],1.2,3.8),
    Pit([5,1],0.9,3.2),
    Pit([15,17],2.2,3.8),
    Pit([12,17],1.2,2.5),
    Pit([7,17],2.0,5.0),
    Pit([1,17],1.8,3.6),
  ]),
  Trial(Date(2019, 12, 5), 10, 0, [17, 16], [
    Pit([17,4],1.3,3.1),
    Pit([10,4],1.5,3.1),
    Pit([16,9],1.4,2.9),
  ]),
  Trial(Date(2019, 12, 19), 12, 4, [8,7], [
    Pit([4,7],.8,.9),
    Pit([3,5],.9,.8),
    Pit([8,2],2,3),
  ]),
  Trial(Date(2019, 12, 20), 5, 0, [8,7], [
    Pit([6,7],.8,.8),
    Pit([2,2],.8,.8),
    Pit([8,6],.8,.8),
    Pit([2,9],.8,.8),
  ]),
]

from sys import argv
from math import sqrt
from numpy import poly1d, polyfit
from scipy.stats import pearsonr
from collections import Counter

if len(argv) < 2:
  print('You must provide at least one argument to choose the function of this program: `img` or `nei` or `depwid')
  quit()

arg = argv[1]

if arg == 'img':
  for trial in trials:
    trial.plot(save=True)
elif arg == 'nei':
  x = []
  y = []
  for trial in trials:
    size = sqrt(trial.size[0]*trial.size[1])
    nei = trial.nearest_neighbor()
    #for n in range(len(nei)):
    #  nei[n] += n*.05
    #x += [sqrt(trial.size[0]*trial.size[1])]*len(nei)
    #y += nei
    x.append(size)
    y.append(sum(nei)/len(nei))
  fig = plt.figure()
  ax = fig.add_subplot(111)
  plt.text(0.1, 0.9, 'R^2 = %.3f\np=%.3f' % (pearsonr(x,y)[0]**2, pearsonr(x,y)[1]), ha='center', va='center', transform=ax.transAxes)
  plt.xlabel('Square root of Trial Area (cm)')
  plt.ylabel('Nearest Neighbor for Individual Pits (cm)')
  plt.plot(x, y, 'bo')
  plt.plot(x, poly1d(polyfit(x, y, 1))(x))
  plt.savefig('nearest_neighbor.png', bbox_inches='tight')
elif arg == 'depwid':
  depths, widths, sizes = [], [], []
  for trial in trials:
    size = sqrt(trial.size[0]*trial.size[1])
    for pit in trial.pits:
      sizes.append(size)
      depths.append(pit.depth)
      widths.append(pit.diam)
  weights = [20*i for i in Counter(depths).values() for j in range(i)]
  plt.scatter([size-.5 for size in sizes], depths, weights, 'b','o', label='depth')
  weights = [20*i for i in Counter(widths).values() for j in range(i)]
  plt.scatter(sizes, widths, weights, 'r', 'o', label='width')
  plt.xlabel('Square root of Trial Area (cm)')
  plt.ylabel('Depth/Width of Antlion Pits (cm)')
  plt.legend(loc='upper right')
  plt.text(10, 7, 'R^2 = %.3f\np=%.3f' % (pearsonr(sizes,depths)[0], pearsonr(sizes,depths)[1]), ha='center', va='center')
  plt.savefig('depth_width.png', bbox_inches='tight')
if arg == 'table':
  print('\\table{')
  print('Dimensions (in)& Pit Depth (cm)& Pit Width (cm)& Nearest Neighbor (cm)')
  for trial in trials:
    size = trial.size
    nei = trial.nearest_neighbor()
    for pitind in range(len(trial.pits)):
      pit = trial.pits[pitind]
      print('& '.join(['x'.join([str(el) for el in size]), '%.1f' % pit.depth, '%.1f' % pit.diam, '%.2f' % nei[pitind]])+'\\cr\\noalign{\\hrule}')
  print('}')
if arg == 'deathtable':
  print('\\table{')
  print('Trial Size& Date& Introduced& Deaths& Pits formed\\cr\\noalign{\\hrule}')
  for trial in trials:
    print('& '.join(['x'.join([str(el) for el in trial.size]), str(trial.date), str(trial.intro), str(trial.dead), str(len(trial.pits))])+'\\cr\\noalign{\\hrule}')
  print('}')