from __future__ import print_function, unicode_literals, division, absolute_import
import math
import numpy

def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / numpy.linalg.norm(vector)

def angle_between(v1, v2):
    """ Returns the angle in radians between vectors 'v1' and 'v2'::

            >>> angle_between((1, 0, 0), (0, 1, 0))
            1.5707963267948966
            >>> angle_between((1, 0, 0), (1, 0, 0))
            0.0
            >>> angle_between((1, 0, 0), (-1, 0, 0))
            3.141592653589793
    """
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    angle = numpy.arccos(numpy.dot(v1_u, v2_u))
    if numpy.isnan(angle):
        if (v1_u == v2_u).all():
            return 0.0
        else:
            return numpy.pi
    return angle




class TrajectoryLeg:
    def __init__(self, thrustvector=None, t=0):
        if thrustvector is None:
            thrustvector = numpy.zeros(3)
        self.thrustvector = thrustvector  #thrust vector
        self.t = t  #length of time to apply the thrust

    def __repr__(self):
        return repr( (self.thrustvector, self.t) )

#
 #* Computes the optimal trajectory from A to B in a solar system using classical
 #* mechanics. Finds the trajectory legs (each expressed as a thrust vector and a
 #* length of time) that as soon as possible will move the ship to the same
 #* position and velocity as the destination.
 #* <P>
 #* The objective is to plan the space trajectory needed to rendezvous with a
 #* given destination within the solar system. It is needed to implement the
 #* game's autopilot which is used by both computer-controlled ships and players.
 #* The trajectory must be computed in advance since simply going in the
 #* direction of the destination would result in arriving with a speed very
 #* different from that of the destination object. We need to know beforehand
 #* when to 'gas' and when to 'break', and in what precise directions to do this.
 #* <P>
 #* The planned trajectory should be fairly efficient (primarily regarding travel
 #* time, but also regarding fuel consumption) and preferably not cheat by using
 #* different physics than the player experiences with the manual controls. It
 #* must also be fairly fast to compute - it's an MMO game and there can be many
 #* thousands of traveling ships.
 #* <P>
 #* The formal problem is: The traveling spaceship, at current position P in
 #* space and with current velocity V, is to travel and rendezvous with a
 #* destination object (a planet, another ship etc) which is at current position
 #* Q and with current velocity W. Since it's a rendezvous (e.g. docking or
 #* landing) the traveling spaceship must have the same velocity W as the
 #* destination upon arrival.
 #* <P>
 #* <I>Input:</I> An instance of TrajectoryParameters that contains:<BR>
 #* deltaPosition<BR>
 #* deltaVelocity<BR>
 #* maxForce (ship's max thrust)<BR>
 #* mass (ship's mass)
 #* <P>
 #* <I>Output:</I> An array of TrajectoryLeg instances, each containing:<BR>
 #* thrust vector (length of which is the thrust force; less than or equal to maxForce)<BR>
 #* time (length of time to apply the thrust)
 #* <P>
 #* The problem solved here makes no mention of changes to the destination's
 #* velocity over time. It's been simplified to disregard the destination's own
 #* acceleration. In theory the change in velocity direction over time is
 #* possible to compute for orbiting celestial bodies, although difficult within
 #* the context of this problem. But other spaceships are impossible to predict
 #* because someone else is controlling them! <I>Therefore the trajectory must be
 #* recomputed regularly during the journey - and we might as well disregard the
 #* complication of the destination's continuously changing acceleration in this
 #* solution.</I>
 #*
 #* @author Christer Swahn
 # (ported to Python anda adapted by proycon)
 #*/


class TrajectoryPlanner:
    #the time tolerance
    TIME_TOLERANCE = 1e-6;
    #defines what is regarded as equivalent position and speed */
    GEOMETRIC_TOLERANCE = 1e-6;
    #the force tolerance during calculation
    FORCE_TOLERANCE = 1e-6;
    #the force difference tolerance of the result
    RESULT_FORCE_TOLERANCE = 1e-2;

    ITER_LIMIT = 30




    def __init__(self, originpos, destpos, originv, destv, maxforce, mass):
        #originpos and destpos are vectors (coordinates) at time 0, same for originv and destv


        self.originpos = originpos
        self.originv = originv

        #we translate to a relative frame of reference
        self.deltapos = destpos - originpos
        self.deltav = destv - originv

        self.maxforce = maxforce
        self.mass = mass


        self.trajectory = [TrajectoryLeg(), TrajectoryLeg()] #will consist of two trajectory legs, each trajectory leg is a (thrustvector,time) tuple


        dV_magnitude = numpy.linalg.norm(self.deltav)
        distance = numpy.linalg.norm(self.deltapos)

        #check special case 1, dV = 0:
        if (dV_magnitude < self.GEOMETRIC_TOLERANCE):
            if (distance < self.GEOMETRIC_TOLERANCE):
                raise Exception("Already at destination")

            t_tot =  (2*self.mass*distance / self.maxforce)**0.5
            self.trajectory[0].thrustvector = self.deltapos * (self.maxforce/distance) #accelerating
            self.trajectory[1].thrustvector =  self.trajectory[0].thrustvector * -1  #breaking
            self.trajectory[0].t = self.trajectory[1].t = t_tot / 2;
            return

        F_v = numpy.zeros(3)
        D_ttot = numpy.zeros(3)
        V_d_ttot = numpy.zeros(3)
        R_d = numpy.zeros(3)
        F_d = numpy.zeros(3)
        F_d2 = numpy.zeros(3)

        #pick f_v in (0, tp.maxForce) via f_v_ratio in (0, 1):
        best_f_v_ratio = -1.0;
        min_f_v_ratio = 0.0;
        max_f_v_ratio = 1.0;
        simple_f_v_ratio = self.calcsimpleratio()
        f_v_ratio = simple_f_v_ratio  #start value
        min_f_v_ratio = simple_f_v_ratio * .99  #(account for rounding error)
        nofIters = 0
        while nofIters < self.ITER_LIMIT:
            nofIters += 1
            f_v = f_v_ratio * self.maxforce;

            t_tot = self.mass / f_v * dV_magnitude

            F_v = self.deltav * (self.mass / t_tot)
            D_ttot = self.deltav * (t_tot/2) + self.deltapos

            dist_ttot = numpy.linalg.norm(D_ttot)

            #check special case 2, dP_ttot = 0:
            if dist_ttot < self.GEOMETRIC_TOLERANCE:
                # done!  F1 = F2 = Fv (only one leg)  (such exact alignment of dV and dP is rare)
                # FUTURE: should we attempt to find more optimal trajectory in case f_v < maxForce?
                if f_v_ratio < 0.5:
                    print("Non-optimal trajectory for special case: dV and dP aligned: f_v_ratio=" + f_v_ratio)
                self.trajectory = [ TrajectoryLeg(F_v, t_tot) ]
                return

            V_d_ttot = D_ttot * (2/t_tot)
            R_d = D_ttot * (1/dist_ttot)  #normalized D_ttot




            alpha = math.pi - angle_between(F_v,R_d)  #angle between F_v and F_p1
            assert (alpha >= 0 and alpha <= math.pi)  #alpha + " not in (0, PI), F_v.dot(R_p1)=" + F_v.dot(R_d);

            if math.pi - alpha < 0.00001:
                #special case 3a, F_v and F_p1 are parallel in same direction
                f_d = self.maxforce - f_v
            elif alpha < 0.00001:
                #special case 3b, F_v and F_p1 are parallel in opposite directions
                f_d = self.maxforce + f_v
            else:
                sin_alpha = math.sin(alpha)
                f_d = self.maxforce/sin_alpha * math.sin(math.pi - alpha - math.asin(f_v/self.maxforce*sin_alpha))
                assert (f_d > 0 and f_d < 2*self.maxforce) #: f_d + " not in (0, " + (2*tp.maxForce) + ")";

            t_1 = 2*self.mass*dist_ttot / (t_tot*f_d)
            t_2 = t_tot - t_1
            if t_2 < self.TIME_TOLERANCE:
                #pick smaller f_v
                #LOG.debug(String.format("Iteration %2d:             f_v_ratio %f; t_2 %,7.3f", nofIters, f_v_ratio, t_2));
                max_f_v_ratio = f_v_ratio
                f_v_ratio += (min_f_v_ratio - f_v_ratio) / 2  #(divisor experimentally calibrated)
                continue

            F_d = R_d * f_d
            F_d2 = F_d * (-t_1/t_2) #since I_d = -I_d2
            #assert (F_d.copy().scale( t_1/tp.mass).differenceMagnitude(V_d_ttot) < GEOMETRIC_TOLERANCE) : F_d;
            #assert (F_d2.copy().scale(-t_2/tp.mass).differenceMagnitude(V_d_ttot) < GEOMETRIC_TOLERANCE) : F_d2;

            F_1 = F_d + F_v
            F_2 = F_d2 + F_v
            #assert (Math.abs(F_1.length()-tp.maxForce)/tp.maxForce < FORCE_TOLERANCE) : "f1=" + F_1.length() + " != f=" + tp.maxForce;
            #assert verifyF1F2(tp, t_1, t_2, F_1, F_2) : "F1=" + F_1 + "; F2=" + F_2;

            f_2 = numpy.linalg.norm(F_2)
            if f_2 > self.maxforce:
                #// pick smaller f_v
                #//LOG.debug(String.format("Iteration %2d:             f_v_ratio %f; f_2 diff %e", nofIters, f_v_ratio, (tp.maxForce-f_2)));
                max_f_v_ratio = f_v_ratio
                f_v_ratio += (min_f_v_ratio - f_v_ratio) / 1.25 #divisor experimentally calibrated)
            else:
                #best so far
                #LOG.debug(String.format("Iteration %2d: best so far f_v_ratio %f; f_2 diff %e; max_f_v_ratio %f", nofIters, f_v_ratio, (tp.maxForce-f_2), max_f_v_ratio));
                best_f_v_ratio = f_v_ratio
                self.trajectory[0].thrustvector = F_1
                self.trajectory[0].t =  t_1
                self.trajectory[1].thrustvector = F_2
                self.trajectory[1].t = t_2
                if f_2 < (self.maxforce*(1-self.RESULT_FORCE_TOLERANCE)):
                    #// pick greater f_v
                    min_f_v_ratio = f_v_ratio;
                    f_v_ratio += (max_f_v_ratio - f_v_ratio) / 4   #// (divisor experimentally calibrated)
                else:
                    break #done!


        if best_f_v_ratio >= 0:
            return
        else:
            print("Couldn't determine full trajectory for %s (nofIters %d) best_f_v_ratio=%.12f" % (tp, nofIters, best_f_v_ratio))
            self.trajectory = [ TrajectoryLeg(self.deltav, dV_magnitude*self.mass/self.maxforce) ]
            #trajectory[0].thrust.add(tp.deltaPosition).normalize().scale(tp.maxForce);  // set thrust direction to average of dP and dV

            # set thrust direction to average of dP and dV
            self.trajectory[0].thrustvector += self.deltapos
            self.trajectory[0].thrustvector = (self.trajectory[0].thrustvector / numpy.linalg.norm(self.trajectory[0].thrustvector)) * self.maxforce #normalize and extend

            return

    #private static boolean verifyF1F2(TrajectoryParameters tp, double t_1, double t_2, SpatialVector F_1, SpatialVector F_2) {
    #    final double REL_TOLERANCE = 1e-5;

    #    SpatialVector V_1 = F_1.copy().scale(t_1/tp.mass);
    #    SpatialVector V_2 = F_2.copy().scale(t_2/tp.mass);
    #    SpatialVector achievedDeltaVelocity = V_1.copy().add(V_2);
    #    double velDiff = achievedDeltaVelocity.differenceMagnitude(tp.deltaVelocity);
    #    assert (velDiff/tp.deltaVelocity.length() < REL_TOLERANCE) : velDiff/tp.deltaVelocity.length() +
    #        "; difference=" + velDiff + "; achievedDeltaVelocity=" + achievedDeltaVelocity + "; tp.deltaVelocity=" + tp.deltaVelocity;

    #    SpatialVector targetPosition = tp.deltaVelocity.copy().scale(t_1+t_2).add(tp.deltaPosition);
    #    SpatialVector achievedPosition = V_1.scale(t_1/2+t_2).add(V_2.scale(t_2/2));
    #    double posDiff = achievedPosition.differenceMagnitude(targetPosition);
    #    assert (posDiff/tp.deltaPosition.length() < REL_TOLERANCE) : posDiff/tp.deltaPosition.length() +
    #        "; difference=" + posDiff + "; achievedPosition=" + achievedPosition + "; targetPosition=" + targetPosition;

    #    return true;
    #}


    def calcsimpleratio(self):
        # compute the f_v ratio of the simplest trajectory where we accelerate in two steps:
        # 1) reduce the velocity difference with the destination to 0 (time needed: t_v)
        # 2) traverse the distance to arrive at the destination (time needed: t_p)
        # (This usually takes longer than the optimal trajectory, since the distance traversal
        # does not utilize the full travel time period. (The greater travel period that
        # the distance traversal can use, the less impulse (acceleration) it needs.))

        inv_acc = self.mass / self.maxforce
        t_v = inv_acc * numpy.linalg.norm(self.deltav)

        newdeltapos = (self.deltav * (t_v/2)) + self.deltapos
        distance = numpy.linalg.norm(newdeltapos)
        t_p = 2 *  (inv_acc*distance)**0.5

        t_tot = t_v + t_p
        f_v_ratio = t_v / t_tot
        return f_v_ratio


    def velocity(self,t,m):
        """Compute velocity vector at time t"""

        pos = self.originpos.copy()
        v = self.originv.copy()

        for trajectoryleg in self.trajectory:
            a = (trajectoryleg.thrustvector) / m  # a = F/m
            if t <= trajectoryleg.t:
                v += a * t
                break
            else:
                v += a * trajectoryleg.t
            t = t - trajectoryleg.t
        return v



#Visualisation stuff



AU = 149597870700 #m

originpos = numpy.array([1*AU,1*AU,1*AU])
originv = numpy.array((-800000,0,0))
destpos = numpy.array([3*AU,0,0])
destv = numpy.array((-50000,-1000000,50000))
mass = 10000
maxforce = mass*11


tp = TrajectoryPlanner(originpos, destpos, originv, destv, maxforce, mass)
print(tp.trajectory)


from visual import * #requires vpython
def np2v(a):
    """Numpy to vpython"""
    return vector(a[0],a[1],a[2])

scene.center = (0,0,0)
scene.width = 1600
scene.height = 1200
scene.range = (5*AU,5*AU,5*AU)

A = sphere(pos=originpos,radius=0.05*AU,color=color.red)
B = sphere(pos=destpos,radius=0.05*AU,color=color.blue)

x = sphere(pos=originpos,radius=0.03*AU,color=color.white,make_trail=True)

t = 0
while True:
    t += 1
    rate(10000) # a high number pretty much means go as fast as our computer can

    A.pos = A.pos + vector(originv)
    B.pos = B.pos + vector(destv)
    v = tp.velocity(t,mass)
    x.pos = x.pos + np2v(v)
    #print(t,v)