from matplotlib.pyplot import plot
from pylab import figure, show, rand
from matplotlib.patches import Ellipse
import numpy as np

def dist(a,b):
    c=a-b
    return np.sqrt(c*c.conjugate())

def n2coord(n):
    return [n.real, n.imag]

coeffs = input('p(z) = az^3 + bz^2 + cz + d \nInput a,b,c,d: ') 
t_vertices = np.roots(coeffs)

#print 'Roots are ' + t_vertices

m = (t_vertices[0] + t_vertices[1])/2.0 #Midpoint of one of the sides

d_coeffs = [3*coeffs[0], 2*coeffs[1], coeffs[2]]
foci = np.roots(d_coeffs)
width = dist(m,foci[0]) +dist(m,foci[1])
center = foci[0] + foci[1]
center = center/2
f = foci[1]/2.0 - foci[0]/2.0
a = width/2.0
height = 2*np.sqrt(a**2 - dist(f,0)**2)

#print 'Foci of ellpise: ' + foci
#print 'Major axis length: ' + width

fig = figure()
ax = fig.add_subplot(111)
e = Ellipse(xy=n2coord(center), width = width, height = height, angle = np.angle(f,deg=True))
ax.add_artist(e)

e.set_alpha(0.5)
e.set_facecolor('blue')


plot([t_vertices[0].real, t_vertices[1].real], [t_vertices[0].imag, t_vertices[1].imag], 'r-')
plot([t_vertices[0].real, t_vertices[2].real], [t_vertices[0].imag, t_vertices[2].imag], 'r-')
plot([t_vertices[1].real, t_vertices[2].real], [t_vertices[1].imag, t_vertices[2].imag], 'r-')

plot([foci[0].real, foci[1].real], [foci[0].imag, foci[1].imag], 'go')

show()