"""Finds out the ellipse that best fits to a set of data points and calculates
   its keplerian elements.
"""

import math
import argparse
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize
from functools import partial

def __read_args():
    """Reads command line arguments.

       Returns:
           object: Parsed arguments.
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--file', type=str, help='path to .csv file', default='../example_data/orbit.csv')
    parser.add_argument('-u', '--units', type=str, help='units of distance (m or km)', default='km')
    return parser.parse_args()


def __cross_sum(data):
    """Returns the normalized sum of the cross products between consecutive vectors.

    Args:
        data(nx3 numpy array): A matrix where each column represents the x,y,z coordinates of each position vector.

    Returns:
        float: The normalized sum of the cross products between consecutive vectors.
    """

    cross_sum = 0
    for i in range(len(data)-1):
        v1 = data[i]
        v2 = data[i+1]
        cross_sum = cross_sum + np.cross(v1,v2)

    return cross_sum/np.linalg.norm(cross_sum)


def __plane_err(data,coeffs):
    """Calculates the total squared error of the data wrt a plane.

       The data should be a list of points. coeffs is an array of
       3 elements - the coefficients a,b,c in the plane equation
       ax+by+c = 0.

       Args:
           data(nx3 numpy array): A numpy array of points.
           coeffs(1x3 array): The coefficients of the plane ax+by+c=0.

       Returns:
           float: The total squared error wrt the plane defined by ax+by+cz = 0.
    """

    a,b,c = coeffs
    return np.sum((a*data[:,0]+b*data[:,1]+c*data[:,2])**2)/(a**2+b**2+c**2)


def __project_to_plane(points,coeffs):
    """Projects points onto a plane.

       Projects a list of points onto the plane ax+by+c=0,
       where a,b,c are elements of coeffs.

       Args:
           points(nx3 numpy array): A numpy array of points.
           coeffs(1x3 array): The coefficients of the plane ax+by+c=0.

       Returns:
           nx3 numpy array: A list of projected points.
    """

    a,b,c = coeffs

    proj_mat =  [[b**2+c**2,  -a*b   ,   -a*c  ],
                 [   -a*b  ,a**2+c**2,   -b*c  ],
                 [   -a*c  ,  -b*c   ,a**2+b**2]]

    return np.matmul(points,proj_mat)/(a**2+b**2+c**2)


def __conv_to_2D(points,x,y):
    """Finds coordinates of points in a plane wrt a basis.

       Given a list of points in a plane, and a basis of the plane,
       this function returns the coordinates of those points
       wrt this basis.

       Args:
           points(numpy array): A numpy array of points.
           x(3x1 numpy array): One vector of the basis.
           y(3x1 numpy array): Another vector of the basis.

       Returns:
           nx2 numpy array: Coordinates of the points wrt the basis [x,y].
    """

    mat = [x[0:2],y[0:2]]
    mat_inv = np.linalg.inv(mat)
    coords = np.matmul(points[:,0:2],mat_inv)

    return coords

def __cart_to_pol(points):
    """Converts a list of cartesian coordinates into polar ones.

       Args:
           points(nx2 numpy array): The list of points in the format [x,y].

       Returns:
           nx2 numpy array: A list of polar coordinates in the format [radius,angle].
    """

    pol = np.empty(points.shape)
    pol[:,0] = np.sqrt(points[:,0]**2+points[:,1]**2)
    pol[:,1] = np.arctan2(points[:,1],points[:,0])

    return pol

def __ellipse_err(polar_coords,params):
    """Calculates the total squared error of the data wrt an ellipse.

       params is a 3 element array used to define an ellipse.
       It contains 3 elements a,e, and t0.

       a is the semi-major axis
       e is the eccentricity
       t0 is the angle of the major axis wrt the x-axis.

       These 3 elements define an ellipse with one focus at origin.
       Equation of the ellipse is r = a(1-e^2)/(1+ecos(t-t0))

       The function calculates r for every theta in the data.
       It then takes the square of the difference and sums it.

       Args:
           polar_coords(nx2 numpy array): A list of polar coordinates in the format [radius,angle].
           params(1x3 numpy array): The array [a,e,t0].

       Returns:
           float: The total squared error of the data wrt the ellipse.
    """

    a,e,t0 = params
    dem = 1+e*np.cos(polar_coords[:,1]-t0)
    num = a*(1-e**2)
    r = np.divide(num,dem)
    err = np.sum((r - polar_coords[:,0])**2)
    return err


def __residuals(data,params,polar_coords,basis):
    """Calculates the residuals after fitting the ellipse.

       Residuals are the difference between the fitted points and
       the actual points.

       res_x = fitted_x - initial_x
       res_y = fitted_y - initial_y
       res_z = fitted_z - initial_z

       where fitted_x,y,z is the closest point on the ellipse to initial_x,y,z.

       However, it is computationally expensive to find the true nearest point.
       So we take an approximation. We consider the point on the ellipse with
       the same true anomaly as the initial point to be the nearest point to it.
       Since the eccentricities of the orbits involved are small, this approximation
       holds.

       Args:
           data(nx3 numpy array): The list of original points.
           params(1x3 numpy array): The array [semi-major axis, eccentricity, argument of periapsis]
                                    of the fitted ellipse.
           polar_coords(nx2 numpy array): The list of 2D polar coordinates of the original points after
                                          projecting them onto the best-fit plane.
           basis(3x2 numpy array): The basis of the best-fit plane.

       Returns:
            nx3 numpy array: Returns the residuals
    """

    a,e,t0 = params
    dem = 1+e*np.cos(polar_coords[:,1]-t0)
    num = a*(1-e**2)
    r = np.divide(num,dem)

    # convert to cartesian
    x_s = np.multiply(r,np.cos(polar_coords[:,1]))
    y_s = np.multiply(r,np.sin(polar_coords[:,1]))

    # convert to 3D
    filtered_coords = np.transpose(np.matmul(basis,[x_s,y_s]))

    residuals = filtered_coords - data

    return residuals

def __read_file(file_name):
    """Reads a space separated csv file with 4 columns in the format t x y z.

       Args:
           file_name(string): the path to the file

       Returns:
           nx3 numpy array: A numpy array with the columns [x y z]. Note that the t coloumn is discarded.
    """

    data = np.loadtxt(file_name,skiprows=1,usecols=(1,2,3))

    return data

def determine_kep(data):
    """Determines keplerian elements that fit a set of points.

       Args:
           data(nx3 numpy array): A numpy array of points in the format [x y z].

       Returns:
           (kep,res) - The keplerian elements and the residuals as a tuple.
           kep: 1x6 numpy array
           res: nx3 numpy array

           For the keplerian elements:
           kep[0] - semi-major axis (in whatever units the data was provided in)
           kep[1] - eccentricity
           kep[2] - inclination (in degrees)
           kep[3] - argument of periapsis (in degrees)
           kep[4] - right ascension of ascending node (in degrees)
           kep[5] - true anomaly of the first row in the data (in degrees)

           For the residuals: (in whatever units the data was provided in)
           res[0] - residuals in x axis
           res[1] - residuals in y axis
           res[2] - residuals in z axis
    """

    # try to fit a plane to the data first.

    # make a partial function of plane_err by supplying the data
    plane_err_data = partial(__plane_err,data)

    # plane is defined by ax+by+cz=0.
    p0 = __cross_sum(data) # make an initial guess

    # minimize the error
    p = minimize(plane_err_data,p0,method='nelder-mead',options={'maxiter':1000}).x
    p = p/np.linalg.norm(p) # normalize p

    # now p is the normal vector of the best-fit plane.

    # lan_vec is a vector along the line of intersection of the plane
    # and the x-y plane.
    lan_vec = np.cross([0,0,1],p)

    # if lan_vec is [0,0,0] it means that it is undefined and can take on
    # any value. So we set it to [1,0,0] so that the rest of the
    # calculation can proceed.
    if (np.array_equal(lan_vec,[0,0,0])):
        lan_vec = [1,0,0]

    # inclination is the angle between p and the z axis.
    inc = math.acos(np.clip(p[2]/np.linalg.norm(p),-1,1))
    # lan is the angle between the lan_vec and the x axis.
    lan = math.atan2(lan_vec[1],lan_vec[0])%(2*math.pi)

    # now we try to convert the problem into a 2D problem.

    # project all the points onto the plane.
    proj_data = __project_to_plane(data,p)

    # p_x and p_y are 2 orthogonal unit vectors on the plane.
    p_x,p_y = lan_vec, np.cross(p,lan_vec)
    p_x,p_y = p_x/np.linalg.norm(p_x), p_y/np.linalg.norm(p_y)

    # find coordinates of the points wrt the basis [p_x,p_y].
    coords_2D = __conv_to_2D(proj_data,p_x,p_y)

    # now try to fit an ellipse to these points.

    # convert them into polar coordinates
    polar_coords = __cart_to_pol(coords_2D)

    # make an initial guess for the parametres
    r_m = np.min(polar_coords[:,0])
    r_M = np.max(polar_coords[:,0])
    a0 = (r_m+r_M)/2
    e0 = (r_M-r_m)/(r_M+r_m)
    t00 = polar_coords[np.argmin(polar_coords[:,0]),1]

    params0 = [a0,e0,t00] # initial guess
    # make a partial function of ellipse_err with the data
    ellipse_err_data = partial(__ellipse_err,polar_coords)
    # minimize the error
    params = minimize(ellipse_err_data,params0,method='nelder-mead',options={'maxiter':1000}).x
    params[2] = params[2]%(2*math.pi)  # bring argp between 0-360 degrees

    # calculate the true anomaly of the first entry in the dataset
    true_anom = (polar_coords[0][1]-params[2])%(2*math.pi)

    # calculation of residuals
    res = __residuals(data,params,polar_coords,np.column_stack((p_x,p_y)))

    kep = np.empty((6,1))
    kep[0] = params[0]
    kep[1] = params[1]
    kep[2] = math.degrees(inc)
    kep[3] = math.degrees(params[2])
    kep[4] = math.degrees(lan)
    kep[5] = math.degrees(true_anom)

    return kep,res

def __print_kep(kep,res,unit):
    """Prints the keplerian elements and some information on residuals.

       Args:
           kep(1x6 numpy array): keplerian elements
           res(nx3 numpy array): residuals
           unit(string): units of distance used

       Returns:
           NIL
    """

    # output the parameters
    print("Semi-major axis:            ",kep[0][0],unit)
    print("Eccentricity:               ",kep[1][0])
    print("Inclination:                ",kep[2][0],"deg")
    print("Argument of periapsis:      ",kep[3][0],"deg")
    print("Longitude of Ascending Node:",kep[4][0],"deg")
    print("True Anomaly                ",kep[5][0],"deg")

    # print data about residuals
    print()

    max_res = np.max(res,axis=0)
    min_res = np.min(res,axis=0)
    sum_res = np.sum(res,axis=0)
    avg_res = np.average(res,axis=0)
    std_res = np.std(res,axis=0)

    print("Printing data about residuals in each axis:")
    print("Max:               ",max_res)
    print("Min:               ",min_res)
    print("Sum:               ",sum_res)
    print("Average:           ",avg_res)
    print("Standard Deviation:",std_res)

def plot_kep(kep,data):
    """Plots the original data and the orbit defined by the keplerian elements.

       Args:
           kep(1x6 numpy array): keplerian elements
           data(nx3 numpy array): original data

       Returns:
           nothing
    """

    a = kep[0]
    e = kep[1]
    inc = math.radians(kep[2])
    t0 = math.radians(kep[3])
    lan = math.radians(kep[4])

    p_x = np.array([math.cos(lan), math.sin(lan), 0])
    p_y = np.array([-math.sin(lan)*math.cos(inc), math.cos(lan)*math.cos(inc), math.sin(inc)])

    # generate 1000 points on the ellipse
    theta = np.linspace(0,2*math.pi,1000)
    radii = a*(1-e**2)/(1+e*np.cos(theta-t0))

    # convert to cartesian
    x_s = np.multiply(radii,np.cos(theta))
    y_s = np.multiply(radii,np.sin(theta))

    # convert to 3D
    mat = np.column_stack((p_x,p_y))
    coords_3D = np.matmul(mat,[x_s,y_s])

    fig = plt.figure()
    ax = Axes3D(fig)
    ax.axis('equal')

    # plot
    ax.plot3D(coords_3D[0],coords_3D[1],coords_3D[2],c = 'red',label='Fitted Ellipse')
    ax.scatter3D(data[:,0],data[:,1],data[:,2],c='black',label='Initial Data')

    # The Pale Blue Dot
    ax.scatter3D(0,0,0,c='blue',depthshade=False,label='Earth')

    ax.can_zoom()
    ax.legend()
    plt.show()

if __name__ == "__main__":
    args = __read_args()
    data = __read_file(args.file)
    kep, res = determine_kep(data)
    __print_kep(kep,res,args.units)
    plot_kep(kep,data)
