from data import trials
from math import sqrt
from scipy.stats import pearsonr
from collections import Counter
import matplotlib.pyplot as plt

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')