from time import sleep

import numpy as np
import numpy.linalg as la

from simple_pid import PID
import krpc


conn = krpc.connect()
print(conn.krpc.get_status().version)

START_ALIGNMENT_SPEED = 900


def get_surface_prograde_pitch(sc, vessel):
    import math

    def cross(u, v):
        return (
            u[1] * v[2] - u[2] * v[1],
            u[2] * v[0] - u[0] * v[2],
            u[0] * v[1] - u[1] * v[0])

    def dot(u, v):
        return u[0] * v[0] + u[1] * v[1] + u[2] * v[2]

    # these vectors are all in vessel.surface_reference_frame
    north = (0, 1, 0)
    east = (0, 0, 1)

    surface_prograde = sc.transform_direction(
        (0, 1, 0),
        vessel.surface_velocity_reference_frame,
        vessel.surface_reference_frame)

    plane_normal = cross(north, east)

    return math.asin(dot(plane_normal, surface_prograde)) * (180.0 / math.pi)


def lift_off(conn, altitude=100000):
    sc = conn.space_center
    mj = conn.mech_jeb
    ascent = mj.ascent_autopilot

    ascent.desired_orbit_altitude = altitude
    ascent.desired_inclination = 0

    ascent.force_roll = True
    ascent.vertical_roll = 0
    ascent.turn_roll = 0

    ascent.autostage = False
    ascent.skip_circularization = True
    ascent.enabled = True
    sc.active_vessel.control.activate_next_stage()

    print("We have lift off!")


def orbit(conn):
    sc = conn.space_center
    mj = conn.mech_jeb
    ascent = mj.ascent_autopilot
    planner = mj.maneuver_planner

    v = sc.active_vessel

    with conn.stream(getattr, ascent, "status") as status:
        status.rate = 1
        with status.condition:
            while status() != 'Off':
                print(status())
                status.wait()

    v.control.rcs = True
    v.control.set_action_group(1, True)

    executor = mj.node_executor
    executor.tolerance = 1

    circ = planner.operation_circularize
    circ.time_selector.time_reference = mj.TimeReference.apoapsis
    circ.make_nodes()
    executor.execute_all_nodes()

    with conn.stream(getattr, executor, "enabled") as enabled:
        enabled.rate = 1
        with enabled.condition:
            while enabled():
                enabled.wait()

    print("Orbit complete!")


def dropping_payload(conn):
    print("Dropping payload")
    sc = conn.space_center
    v = sc.active_vessel
    ctrl = v.control

    payload = ctrl.activate_next_stage()
    payload = payload[0]

    ctrl.up = -1
    with conn.stream(payload.position, v.reference_frame) as payload_position:
        payload_position.rate = 1
        far_enough = False
        while not far_enough:
            displacement = np.array(payload_position())
            distance = la.norm(displacement)
            far_enough = distance > 20

    ctrl.up = 0
    print("Payload Dropped")


def is_longitude_in_target_range(long, target, range):
    if target - range <= -180:
        return long < target + range or long > 180 - (-180 - (target - range))
    elif target + range > 180:
        return long > target - range or long < -180 + (180 - (target + range))

    return target - range < long < target + range


TARGET_DEORBIT_NODE_LONGITUDE = 178


def deorbit(conn):
    sc = conn.space_center
    v = sc.active_vessel
    o = v.orbit
    b = o.body

    surface_rf = b.reference_frame

    node = v.control.add_node(sc.ut + 300, -100)

    longitude_at = b.longitude_at_position(o.position_at(node.ut, surface_rf), surface_rf)
    while not is_longitude_in_target_range(longitude_at, TARGET_DEORBIT_NODE_LONGITUDE, 10):
        node.ut = node.ut + 10
        longitude_at = b.longitude_at_position(o.position_at(node.ut, surface_rf), surface_rf)
    while not is_longitude_in_target_range(longitude_at, TARGET_DEORBIT_NODE_LONGITUDE, 1):
        node.ut = node.ut + 1
        longitude_at = b.longitude_at_position(o.position_at(node.ut, surface_rf), surface_rf)

    print("Maneuver created")
    mj = conn.mech_jeb
    executor = mj.node_executor
    executor.tolerance = 1
    executor.execute_one_node()

    with conn.stream(getattr, executor, "enabled") as enabled:
        enabled.rate = 1
        with enabled.condition:
            while enabled():
                enabled.wait()

    v.control.activate_next_stage()
    sc.rails_warp_factor = 7
    print("Ship deorbited")


REENTRY_PITCH_0_SPEED = 1900
REENTRY_PITCH_PROP_SPEED = 600

CHECKPOINTS = [30000, 20000, 10000]


def reentry(conn):
    print("Handling reentry")
    sc = conn.space_center
    v = sc.active_vessel
    mj = conn.mech_jeb
    sa = mj.smart_ass

    sa.autopilot_mode = mj.SmartASSAutopilotMode.surface
    sa.surface_heading = 90
    sa.force_heading = True
    sa.force_pitch = True
    sa.surface_roll = 0
    sa.force_roll = True
    sa.update(True)

    altitude = v.flight().surface_altitude
    old_altitude = altitude
    with conn.stream(v.velocity, v.orbit.body.reference_frame) as velocity:
        velocity.rate = 1
        with velocity.condition:
            speed = la.norm(np.array(velocity()))
            while speed > START_ALIGNMENT_SPEED:
                if speed < REENTRY_PITCH_0_SPEED:
                    if sa.surface_pitch > 0:
                        sa.surface_pitch = 0
                        sa.update(False)
                else:
                    if sa.surface_pitch != 5:
                        sa.surface_pitch = 5
                        sa.update(False)

                altitude = v.flight().surface_altitude
                for point in CHECKPOINTS:
                    if old_altitude > point > altitude:
                        # sc.save(str(point))
                        # sc.quicksave()
                        print("Checkpoint {}: v={} t={}".format(point, speed, v.parts.all[11].skin_temperature))
                        break
                old_altitude = altitude


                velocity.wait()
                speed = la.norm(np.array(velocity()))

    v.control.rcs = False
    print("Reentry completed")


def alignment(conn):
    print("Starting alignment")
    sc = conn.space_center
    v = sc.active_vessel
    b = v.orbit.body
    mj = conn.mech_jeb
    sa = mj.smart_ass
    ref = b.reference_frame

    target_lat = -0.0485997
    factor = 5
    limit_heading = 5
    limit_roll = 3
    pid_h = PID(9, .005, .05, setpoint=target_lat, output_limits=[-limit_heading / factor, limit_heading / factor])

    with open('track.csv', 'w') as track_file:
        while v.flight().surface_altitude > 4000:
            current_lat = b.latitude_at_position(v.position(ref), ref)
            diff = str(target_lat - current_lat).replace('.', ',')
            print(target_lat - current_lat)
            track_file.write("{};\n".format(diff))
            control = pid_h(current_lat)

            correction = control * factor
            sa.surface_heading = 90 - correction
            sa.surface_roll = min(-correction, limit_roll) if correction < 0 else max(-correction, -limit_roll)
            sa.surface_pitch = min(get_surface_prograde_pitch(sc, v) + 7, 0)

            sa.update(False)
            sleep(1)

    v.control.gear = True
    print("Alignment done")


def final_approach(conn):
    print("Final approach")
    sc = conn.space_center
    v = sc.active_vessel
    b = v.orbit.body
    mj = conn.mech_jeb
    sa = mj.smart_ass
    ref = b.reference_frame

    airstrip_position = b.surface_position(-0.0485997, -74.724375, ref)
    airstrip_altitude = b.surface_height(0-.0485997, -74.724375)
    target_vspeed = 40

    h_speed = v.flight(ref).horizontal_speed

    target_lat = -0.0485997
    factor = 5
    limit_heading = 5
    limit_roll = 3
    pid_h = PID(9, .005, .05, setpoint=target_lat, output_limits=[-limit_heading / factor, limit_heading / factor])
    pid_vh = PID(.5, .1, .05, setpoint=h_speed, output_limits=[0, 1])
    while v.flight().surface_altitude > 100:
        ship_position = v.position(ref)
        conn.drawing.add_line(airstrip_position, ship_position, ref)
        distance = np.linalg.norm(np.array(ship_position) - np.array(airstrip_position))
        height = v.flight(ref).mean_altitude - airstrip_altitude

        # angle = math.acos(height/distance)

        current_vspeed = v.flight(v.orbit.body.reference_frame).vertical_speed
        current_hspeed = v.flight(v.orbit.body.reference_frame).horizontal_speed

        # target_hspeed = current_vspeed * math.cos(angle) * -1
        target_hspeed = current_vspeed * (distance / height) * -1
        pid_vh.setpoint = target_hspeed
        control = pid_vh(current_hspeed)

        v.control.throttle = control

        current_lat = b.latitude_at_position(v.position(ref), ref)
        control = pid_h(current_lat)

        correction = control * factor
        sa.surface_heading = 90 - correction
        sa.surface_roll = min(-correction, limit_roll) if correction < 0 else max(-correction, -limit_roll)
        sa.surface_pitch = min(get_surface_prograde_pitch(sc, v) + 7, 0)

        sa.update(False)
        sleep(1)

    # pid_vs = PID(.5, .1, .05, setpoint=5)
    while v.situation == sc.VesselSituation.flying:
        # current_vertical_speed = v.flight(ref).vertical_speed
        # control_pitch = pid_vs(current_vertical_speed)
        sa.surface_heading = 90
        sa.surface_roll = 0
        sa.surface_pitch = 0
        sa.update(False)

        sleep(1)

    v.control.brakes = True
    sa.autopilot_mode = mj.SmartASSAutopilotMode.off

    print("Landed!!")


# check status to determine course of actions
sc = conn.space_center
v = sc.active_vessel
if v.situation == sc.VesselSituation.pre_launch:
    lift_off(conn)

mj = conn.mech_jeb
ascent = mj.ascent_autopilot
if ascent.status and ascent.status != 'Off':
    orbit(conn)

if v.control.current_stage == 3:
    dropping_payload(conn)

if v.control.current_stage == 2 and v.situation == sc.VesselSituation.orbiting:
    deorbit(conn)

current_velocity = la.norm(np.array(v.velocity(v.orbit.body.reference_frame)))
is_in_reentry = v.situation == sc.VesselSituation.sub_orbital \
                or v.situation == sc.VesselSituation.flying and current_velocity > START_ALIGNMENT_SPEED

if v.control.current_stage == 1 and is_in_reentry:
    reentry(conn)

if v.situation == sc.VesselSituation.flying and v.flight().surface_altitude > 4000:
    alignment(conn)

final_approach(conn)

conn.close()