import numpy as np
import logging

class FitEllipse (object):

    def __init__(self,min_points,max_iter,threshold,num_close):

        # points = np.array(candidate_points)
        # y,x = points.T

        C = np.zeros([6,6])
        C[0,2]= 2.0
        C[2,0]= 2.0
        C[1,1]= -1.0

        #self.x = x
        #self.y = y
        self.C = C

        self.min_points = min_points
        self.max_iter = max_iter
        self.threshold = threshold
        self.num_close = num_close

        self.best_params = None
        self.best_params_set = False
        self.besterror = np.inf

    def ransac_fit(self,candidate_points):

        #points = np.array(candidate_points)

        for i in range(self.max_iter):

            inlier_points, outlier_points = self.choose_inliers(candidate_points)
            params, error = self.fit_ellipse(inlier_points)

            if len(outlier_points)>0:
                cost = self.outlier_cost(outlier_points,params)
                also_in = 0
                for j,c in enumerate(cost):
                    point = outlier_points[j]
                    if cost[j]<self.threshold:
                        inlier_points += [point]
                        also_in += 1

                if also_in > self.num_close:
                    params, error = self.fit_ellipse(inlier_points)
                    if (error < self.besterror):
                        self.best_params = params
                        self.best_params_set = True
                        self.besterror = error

        if self.best_params_set:
            return ellipse_center(self.best_params), ellipse_angle_of_rotation(self.best_params)*180./np.pi, ellipse_axis_length(self.best_params)
        else:
            return None

    def choose_inliers(self, candidate_points):

        #cannot take a larger sample than population
        if(len(candidate_points) > self.min_points):
            inlier_index = np.random.choice(np.arange(len(candidate_points)),self.min_points,replace=False)
        else:
            #TODO check this
            inlier_index = np.arange(self.min_points)

        inlier_points = []
        outlier_points = []

        for i in range(len(candidate_points)):
            if i in inlier_index:
                inlier_points += [candidate_points[i]]
            else:
                outlier_points += [candidate_points[i]]

        return inlier_points, outlier_points

    def outlier_cost(self,outlier_points,params):

        y,x = np.array(outlier_points).T

        D = np.vstack([x*x, x*y, y*y, x, y, np.ones(len(y))])
        #S = np.dot(D, D.T)

        cost = (np.dot(params,D))**2

        return cost

    def fit_ellipse(self,inlier_points):
        try:
            inlier_points = np.array(inlier_points)
            points = np.array(inlier_points)
            y,x = points.T

            D = np.vstack([x*x, x*y, y*y, x, y, np.ones(len(y))])
            S = np.dot(D, D.T)

            M = np.dot(np.linalg.inv(S),self.C)
            U,s,V=np.linalg.svd(M)

            params = U.T[0]
            error = np.dot(params, np.dot(S,params))/len(inlier_points)
        except:
            #TODO - check if this is correct
            params = None      #WBW error handling
            error = 0.00000001 #WBW error handling

        return params, error


def ellipse_center(a):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    num = b*b-a*c
    x0=(c*d-b*f)/num
    y0=(a*f-b*d)/num
    return np.array([x0,y0])

def ellipse_angle_of_rotation( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    return 0.5*np.arctan(2*b/(a-c))

def ellipse_angle_of_rotation2( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    if b == 0:
        if a > c:
            return 0
        else:
            return np.pi/2
    else:
        if a > c:
            return np.arctan(2*b/(a-c))/2
        else:
            return np.pi/2 + np.arctan(2*b/(a-c))/2


def ellipse_axis_length( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    up = 2*(a*f*f+c*d*d+g*b*b-2*b*d*f-a*c*g)
    down1=(b*b-a*c)*( (c-a)*np.sqrt(1+4*b*b/((a-c)*(a-c)))-(c+a))
    down2=(b*b-a*c)*( (a-c)*np.sqrt(1+4*b*b/((a-c)*(a-c)))-(c+a))

    #TODO check this - cannot divide by 0 so just use a small number instead
    if(down1 == 0):
        down1 = .0000000001

    if(down2 == 0):
        down2 = .0000000001

    res1=np.sqrt(up/down1)
    res2=np.sqrt(up/down2)
    return np.array([res1, res2])

def fit_ellipse(candidate_points):

    # method from http://nicky.vanforeest.com/misc/fitEllipse/fitEllipse.html
    points = np.array(candidate_points)
    y,x = points.T
    D = np.vstack([x*x, x*y, y*y, x, y, np.ones(len(y))])
    S = np.dot(D, D.T)

    C = np.zeros([6,6])
    C[0,2]= 2.0
    C[2,0]= 2.0
    C[1,1]= -1.0

    M = np.dot(np.linalg.inv(S),C)

    U,s,V=np.linalg.svd(M)

    params = U.T[0]

    return ellipse_center(params), ellipse_angle_of_rotation(params)*180./np.pi, ellipse_axis_length(params)

def rotate_vector(y,x,theta):

    xp = x*np.cos(theta) - y*np.sin(theta)
    yp = x*np.sin(theta) + y*np.cos(theta)

    return yp,xp

def test_fit():

    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    x = np.linspace(-3.0,3.0,1000)
    yp = np.sqrt(4.0 - (4.0/9.0)*(x**2))
    ym = -yp

    yp += 0.1*np.random.normal(size=len(yp))
    ym += 0.1*np.random.normal(size=len(yp))

    yp, x1 = rotate_vector(yp,x,np.pi/8)
    ym, x2 = rotate_vector(ym,x,np.pi/8)

    y_outlier = np.random.random(size=100)*4.0 - 2.0
    x_outlier = np.random.random(size=100)*6.0 - 3.0


    outlier_points = np.vstack([y_outlier, x_outlier]).T

    candidate_points = np.vstack([np.hstack([yp,ym]), np.hstack([x,x])]).T

    candidate_points = np.vstack([outlier_points, candidate_points])

    yt, xt = candidate_points.T
    print xt

    #center, angle, (axis1,axis2) = fit_ellipse(candidate_points)

    fe=FitEllipse(40,100,0.0001,40)
    result = fe.ransac_fit(candidate_points)
    if result!=None:
        center, angle, (axis1,axis2) = fe.ransac_fit(candidate_points)

    print "center = ", center
    print "angle = ", angle
    print "axis1 = ", axis1
    print "axis2 = ", axis2

    fig,ax=plt.subplots(1)
    ax.plot(x,yp,'bo')
    ax.plot(x,ym,'bo')
    ax.plot(x_outlier, y_outlier, 'rx')

    from matplotlib.patches import Ellipse
    el = Ellipse(center,width=2.0*axis1,height=2.0*axis2,angle=angle,fill=False,linewidth=3,color='r')

    ax.add_artist(el)


    plt.show()

if __name__=='__main__':

    test_fit()
