# Data
from scipy.spatial import Voronoi, voronoi_plot_2d, KDTree, distance
from enum import Enum
import matplotlib.pyplot as plt

fancy = {'pits':'Artificial Pits', 'obstacles':'Artificial Obstacles',
'trails':'Trail Erasure'}

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 Trial:
  def __init__(self, method, recl, dead, size, pits, fake=None):
    self.method = method
    # intro = 15
    self.recl = recl
    self.dead = dead
    self.size = size
    self.pits = pits
    self.pitlocs = [pit.loc for pit in self.pits]
    self.fake = fake

  def __repr__(self):
    return f'{fancy[self.method]} {self.size[0]}'

  def plot(self, save=False):
    vor = Voronoi([pit.loc for pit in self.pits])
    voronoi_plot_2d(vor)
    plt.xlabel('%s (dimension %dx%dcm)' % (fancy[self.method], self.size[0], self.size[1]))
    if save:
      plt.savefig('%s-%dx%d.png'%(self.method,self.size[0],self.size[1]),
      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]]

class Obstacle:
    def __init__(self, coords):
        self.coords = coords

trials = [
  Trial('trails', 1, 2, [24,24], [
    Pit((24,1),1,2),
    Pit((21,6),1,1),
    Pit((18,18),2,3),
    Pit((18,24),1,1),
    Pit((23,23),5,3),
    Pit((23,13),2,1),
    Pit((20,6),1,1),
    Pit((4,4),2,3),
    Pit((14,5),4,5),
    Pit((16,9),1,1),
    Pit((9,10),3,4),
    Pit((13,12),1,1),
  ]),
  Trial('trails', 2, 3, [12,12], [
    Pit((12,12),1,2),
    Pit((12,0),1,1),
    Pit((11,8),1,1),
    Pit((11,5),4,7),
    Pit((9,7),2,2),
    Pit((9,9),3,4),
    Pit((5,11),1,3),
    Pit((5,9),1,3),
    Pit((1,11),1,1),
    Pit((3,4),1,3),
  ]),
  Trial('pits', 1, 7, [24,24], [
    Pit((1,23),1,1),
    Pit((3,12),3,5),
    Pit((17,5),2,5),
    Pit((10,17),2,4),
    Pit((23,0),1,1),
    Pit((23,9),5,8),
    Pit((24,24),1,1)
  ], fake=[
    Pit((5,4),5,8),
    Pit((7,2),5,8),
    Pit((20,2),5,8),
    Pit((23,4),5,8),
    Pit((12,7),5,8),
    Pit((7,18),5,8),
    Pit((20,18),5,8),
    Pit((5,17),5,8),
    Pit((12,23),5,8),
    Pit((14,23),5,8),
    Pit((18,23),5,8),
    Pit((14,18),5,8),
  ]),
  Trial('pits', 4, 4, [12,12], [
    Pit((2,12),1,2),
    Pit((6,4),2,3),
    Pit((6,7),5,8),
    Pit((8,12),2,3),
    Pit((12,5),1,1),
    Pit((11,9),1,1),
    Pit((2,9),1,1),
  ], fake=[
    Pit((2,3),5,8),
    Pit((2,12),5,8),
    Pit((6,3),5,8),
    Pit((6,7),5,8),
    Pit((11,3),5,8),
    Pit((11,11),5,8),
  ]),
  Trial('obstacles', 0, 4, [24,24], [
    Pit((1,20),1,1),
    Pit((2,12),1,2),
    Pit((12,0),1,1),
    Pit((12,12),1,2),
    Pit((20,15),5,5),
    Pit((20,24),2,3),
    Pit((11,22),3,5),
    Pit((20,5),3,5),
    Pit((21,8),2,6),
    Pit((20,5),3,5),
    Pit((21,8),2,6),
    Pit((22,14),1,1),
    Pit((22,16),2,4),
  ], fake=[
    Obstacle([(8,3),(1,7),(7,16)]),
    Obstacle([(18,1),(16,7),(21,4),(21,7)]),
    Obstacle([(17,15),(18,20),(23,19),(21,20)])
  ]),
  Trial('obstacles', 3, 0, [12,12], [
    Pit((2,11),2,3),
    Pit((2,1),1,1),
    Pit((5,3),1,3),
    Pit((6,10),1,1),
    Pit((12,4),3,6),
    Pit((11,6),1,3),
    Pit((11,9),1,2),
    Pit((10,11),2,4),
    Pit((6,3),3,5),
    Pit((12,9),1,2),
  ], fake=[ Obstacle([(10,2),(11,9),(3,5),(3,8)]) ]),
]