# bertrand.py
'''
  Betrand's Paradox
  https://en.wikipedia.org/wiki/Bertrand_paradox_(probability)

  Based on code by Christian Hill,
  https://scipython.com/blog/bertrands-paradox/
'''

import math, random
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D


GREY = (0.2,0.2,0.2)
NCHORDS_TO_PLOT = 1000
NUM_CHORDS = 10000
R = 1  # radius
# side length of equilateral triangle inscribed in circle
TRI_LEN = R * math.sqrt(3)


def bertrand1():
  ''' The "random endpoints" method:
    Pairs of (uniformly-distributed) random points on the 
    circumference are joined as chords.
  '''
  chords = []
  midPts = []
  for i in range(NUM_CHORDS):
    angle1 = random.random()*2*math.pi
    angle2 = random.random()*2*math.pi
    chord = ( R*math.cos(angle1), R*math.sin(angle1),
              R*math.cos(angle2), R*math.sin(angle2) )
    chords.append(chord)
    midPt = ((chord[0]+chord[2])/2, (chord[1]+chord[3])/2)
    midPts.append(midPt)
  return "random endpoints", chords, midPts

def bertrand2():
  ''' The "random radial point" method:
    Select a random radius, and then choose a point
    at random (uniformly-distributed) on this radius 
    to be the midpoint of the chosen chord.
  '''
  angles = [ random.random()*2*math.pi 
                             for i in range(NUM_CHORDS)]
  radii =  [ random.random()*R for i in range(NUM_CHORDS)]
  midPts = []
  for i in range(NUM_CHORDS):
    midPt = (radii[i] * math.cos(angles[i]), 
             radii[i] * math.sin(angles[i]))
    midPts.append(midPt)
  chords = getChords(midPts)
  return "random radial point", chords, midPts

def getChords(midPts):
  # Return the chords with the provided midPts
  chords = [(0,0,0,0) for i in range(NUM_CHORDS)]
  for i, (x0, y0) in enumerate(midPts):
    # y = mx + c is the equation of the chord
    m = -x0/y0
    c = y0 + x0**2/y0
    '''
      Solve the quadratic equation where the chord 
      y = mx+c intersects the circle x^2 + y^2 = R^2
      to find two (x,y) endpts.
    '''
    A, B, C = m**2+1, 2*m*c, c**2-R**2
    d = math.sqrt(B**2 - 4*A*C)
    xEnd1 = (-B + d)/ 2 / A
    xEnd2 = (-B - d)/ 2 / A
    yEnd1 = m*xEnd1 + c
    yEnd2 = m*xEnd2 + c
    chords[i] = (xEnd1, yEnd1, xEnd2, yEnd2)
  return chords

def bertrand3():
  ''' The "random midpoint" method:
    Select a point at random (uniformly distributed) 
    within the circle, as the midpoint of the 
    chosed chord.
    Weigh the radial distance by the square root of 
    the random number (0,1] so there is a 
    greater probability for points further from the 
    centre, where there is more room for them.
  '''
  angles = [ random.random()*2*math.pi 
                           for i in range(NUM_CHORDS)]
  radii =  [ math.sqrt(random.random())*R  
                           for i in range(NUM_CHORDS)]
  midPts = []
  for i in range(NUM_CHORDS):
    midPt = (radii[i] * math.cos(angles[i]), 
             radii[i] * math.sin(angles[i]))
    midPts.append(midPt)
  chords = getChords(midPts)
  return "random midpoint", chords, midPts

def createAxes():
  '''
    three methods in three columns;
    a column contains a line drawing above a scatter chart
    each drawn inside a circle
  '''
  fig, axes = plt.subplots(2, 3, figsize=(10, 6),
                    subplot_kw={'aspect': 'equal'})
  for r in range(2):
    for c in range(3):
      circle = Circle((0,0), R, facecolor='none')
      axes[r,c].add_artist(circle)
      axes[r,c].set_xlim((-R,R))
      axes[r,c].set_ylim((-R,R))
      axes[r,c].axis('off')
  return fig, axes

def plotBertrand(fn, axTop, axBot):
  ''' Plot the chords and their midPts on 
    separate axes for the selected function.
  '''
  title, chords, midPts = fn()

  # plot (some of) the chords;
  # keep track of which chords are longer than TRI_LEN
  isLong = [False] * NUM_CHORDS
  for i, chord in enumerate(chords):
    (x0, y0, x1, y1) = chord
    if math.hypot(x0-x1, y0-y1) > TRI_LEN:
      isLong[i] = True
    if i < NCHORDS_TO_PLOT:
      line = Line2D([x0,x1],[y0,y1], color=GREY, alpha=0.1)
      axTop.add_line(line)
  prob = sum(isLong)/NUM_CHORDS
  axTop.set_title(f"{title}; prob: {prob}", fontsize=10)
  print(f"\'{title}\' prob: {prob}")

  # plot the midpoints
  xms, yms = map(list, zip(*midPts))
  axBot.scatter(xms, yms, s=0.2, color=GREY)

fig, axes = createAxes()
plotBertrand(bertrand1, axes[0,0], axes[1,0]) # col 1
plotBertrand(bertrand2, axes[0,1], axes[1,1]) # 2
plotBertrand(bertrand3, axes[0,2], axes[1,2]) # 3
plt.show()