# Bearings only tracking example
#
# 
#
# A submarine moves in a zig-zag pattern to enable tracking of a frigate
# from only angle measurments.
#


import numpy as np
import matplotlib.pyplot as plt

# Time and velocity parameters
dt = 1
vobs = 15
n = 100
# Play with submarine trajectory, 
# set phase to 2 or 20
phase=2
#phase = 6
#phase=20

# Submarine trajectory
msub = np.zeros((n, 2))
for i in range(n - 1):
    angle = (np.pi * 35 / 180) + (np.pi / 2) * np.cos(np.pi * i / phase)
    msub[i + 1, 0] = msub[i, 0] + dt * vobs * np.sin(angle)
    msub[i + 1, 1] = msub[i, 1] + dt * vobs * np.cos(angle)


# Frigate trajectory
vx, vy = -15, -5
mfrig = np.zeros((4, n))
mfrig[:, 0] = [2500, 2500, vx, vy]
A = np.array([[1, 0, dt, 0],
              [0, 1, 0, dt],
              [0, 0, 1, 0],
              [0, 0, 0, 1]])

for i in range(n - 1):
    mfrig[:, i + 1] = A @ mfrig[:, i]
   
    
# Measurement simulation
np.random.seed(10)
sigm = 0.005 ** 2
dat = np.zeros(n)
data = np.zeros(n)
for i in range(n):
    dx = mfrig[0, i] - msub[i, 0]
    dy = mfrig[1, i] - msub[i, 1]
    dat[i] = np.arctan2(dy, dx)
    data[i] = dat[i] + np.sqrt(sigm) * np.random.randn()


# Extended Kalman filter initialization
mprior = np.array([3500, 3500, -10, -10])
sigprior = np.array([
    [1000**2, 0.7*1000*1000, 0.5*1000*4, 0],
    [0.7*1000*1000, 1000**2, 0, 0.5*1000*4],
    [0.5*1000*4, 0, 4**2, 0],
    [0, 0.5*1000*4, 0, 4**2]
])
Sigeps = np.diag([1**2, 1**2, 0.1**2, 0.1**2])

mprop = np.zeros_like(mfrig)
mprop[:, 0] = mprior - np.array([msub[0, 0], msub[0, 1], 0, 0])
mupd = np.zeros_like(mfrig)
sigpred = np.zeros((4, 4, n))
sigpred[:, :, 0] = sigprior
sigupd = np.zeros_like(sigpred)
sserr = np.zeros((2, n))
mekf = np.zeros_like(mfrig)

# EKF tracking

for i in range(n):
    # -------------------------------
    # EKF: Linearization
    # -------------------------------
    denom = 1 + (mprop[1, i] / mprop[0, i]) ** 2
    g = 1 / denom
    gm = g * np.array([
        -mprop[1, i] / (mprop[0, i] ** 2),
        1 / mprop[0, i],
        0, 0
    ])

    # -------------------------------
    # EKF: Update Step
    # -------------------------------
    Smat = gm @ sigpred[:, :, i] @ gm.T + sigm
    Kmat = sigpred[:, :, i] @ gm.T / Smat
    innovation = data[i] - np.arctan2(mprop[1, i], mprop[0, i])
    mupd[:, i] = mprop[:, i] + Kmat * innovation
    sigupd[:, :, i] = (np.eye(4) - np.outer(Kmat, gm)) @ sigpred[:, :, i]

    # -------------------------------
    # For Uncertainty plotting
    # -------------------------------
    sserr[0, i] = np.sqrt(sigupd[0, 0, i])
    sserr[1, i] = np.sqrt(sigupd[1, 1, i])
    mekf[:, i] = mupd[:, i] + np.array([msub[i, 0], msub[i, 1], 0, 0])

    # -------------------------------
    # EKF: Prediction Step
    # -------------------------------
    if i < n - 1:
        delta = msub[i] - msub[i + 1]
        mprop[:, i + 1] = A @ mupd[:, i] + np.array([delta[0], delta[1], 0, 0])
        sigpred[:, :, i + 1] = A @ sigupd[:, :, i] @ A.T + Sigeps

# Visualization
for i in range(n):
    # --- Map View ---
    plt.figure(1)
    plt.clf()
    plt.plot(msub[:i+1, 0], msub[:i+1, 1], 'r-', linewidth=2, label='Submarine')
    plt.plot(mfrig[0, :i+1], mfrig[1, :i+1], 'k-', linewidth=2, label='Frigate')
    plt.plot(mekf[0, :i+1], mekf[1, :i+1], 'b--', label='EKF Estimate')
    plt.plot(mekf[0, i], mekf[1, i], 'bx', markersize=8)

    plt.plot(mekf[0, i] + 2 * sserr[0, i], mekf[1, i] + 2 * sserr[1, i], 'gx')
    plt.plot(mekf[0, i] - 2 * sserr[0, i], mekf[1, i] - 2 * sserr[1, i], 'gx')
    plt.title('Map View')
    plt.xlabel('East')
    plt.ylabel('North')
    plt.axis([-1000, 7000, -1000, 7000])
    plt.legend()


    plt.pause(0.2 if i > 0 else 10)