"""
Physics-embedded CW dynamics for spacecraft terminal rendezvous.

Implements:
- Exact ZOH discretization of CW equations (Phi, Gamma closed-form)
- Continuous-time CW dynamics (A_cw, B)
- CW drift acceleration a_cw(x)
- Discrete LQR nominal controller via DARE
- Kalman filter for state estimation
- Domain randomization utilities

Reference: Section 2 of the paper.
"""

import numpy as np
from scipy.linalg import solve_discrete_are


class CWDynamics:
    """
    Clohessy-Wiltshire relative dynamics with exact zero-order-hold discretization.

    State: x = [r_x, r_y, r_z, v_x, v_y, v_z]^T
    Coordinate: x radial-out, y along-track, z normal

    Parameters
    ----------
    n : float
        Mean orbital angular velocity (rad/s).
    dt : float
        Sampling period (s).
    d_bar : float
        Disturbance acceleration bound (m/s^2).
    """

    def __init__(self, n=None, dt=1.0, d_bar=0.0,
                 mu=3.986004418e14, R0=42164000.0):
        if n is None:
            n = np.sqrt(mu / R0 ** 3)
        self.n = n
        self.dt = dt
        self.d_bar = d_bar
        self.mu = mu
        self.R0 = R0

        self.A_cw = self._build_A_cw(n)
        self.B_ct = self._build_B()
        self.Phi, self.Gamma = self._build_exact_discrete(n, dt)

    @staticmethod
    def _build_A_cw(n):
        """Continuous-time A matrix of CW equations."""
        A = np.zeros((6, 6))
        A[0:3, 3:6] = np.eye(3)
        A[3, 0] = 3.0 * n ** 2
        A[5, 2] = -n ** 2
        A[3, 4] = 2.0 * n
        A[4, 3] = -2.0 * n
        return A

    @staticmethod
    def _build_B():
        B = np.zeros((6, 3))
        B[3:6, :] = np.eye(3)
        return B

    @staticmethod
    def _build_exact_discrete(n, dt):
        """
        Exact ZOH discretization — closed-form Phi and Gamma.
        See Section 2.2 of the paper.
        """
        nu = n * dt
        c = np.cos(nu)
        s = np.sin(nu)

        Phi = np.array([
            [4.0 - 3.0 * c,       0.0, 0.0,
             s / n, 2.0 * (1.0 - c) / n, 0.0],
            [6.0 * (s - nu),       1.0, 0.0,
             -2.0 * (1.0 - c) / n, (4.0 * s - 3.0 * nu) / n, 0.0],
            [0.0, 0.0, c,
             0.0, 0.0, s / n],
            [3.0 * n * s,          0.0, 0.0,
             c, 2.0 * s, 0.0],
            [-6.0 * n * (1.0 - c), 0.0, 0.0,
             -2.0 * s, 4.0 * c - 3.0, 0.0],
            [0.0, 0.0, -n * s,
             0.0, 0.0, c],
        ])

        Gamma = np.array([
            [(1.0 - c) / n ** 2,
             2.0 * (nu - s) / n ** 2, 0.0],
            [-2.0 * (nu - s) / n ** 2,
             (4.0 * (1.0 - c) - 1.5 * nu ** 2) / n ** 2, 0.0],
            [0.0, 0.0, (1.0 - c) / n ** 2],
            [s / n, 2.0 * (1.0 - c) / n, 0.0],
            [-2.0 * (1.0 - c) / n, (4.0 * s - 3.0 * nu) / n, 0.0],
            [0.0, 0.0, s / n],
        ])

        return Phi, Gamma

    def acw(self, x):
        """CW drift acceleration: a_cw(x) from Section 3.4."""
        n = self.n
        return np.array([
            3.0 * n ** 2 * x[0] + 2.0 * n * x[4],
            -2.0 * n * x[3],
            -n ** 2 * x[2],
        ])

    def step(self, x, u, disturbance=None):
        """
        Exact ZOH discrete step: x_{k+1} = Phi x_k + Gamma u_k + w_k
        """
        x_next = self.Phi @ x + self.Gamma @ u
        if disturbance is not None:
            x_next += self.Gamma @ disturbance
        return x_next

    def rollout(self, x0, u_sequence):
        """Nominal multi-step rollout (no disturbance)."""
        L = len(u_sequence)
        traj = np.zeros((L + 1, 6))
        traj[0] = x0
        for k in range(L):
            traj[k + 1] = self.Phi @ traj[k] + self.Gamma @ u_sequence[k]
        return traj

    def sensitivity(self, ell):
        """Proposition 1: d x_{k+ell} / d u_k = Phi^{ell-1} Gamma"""
        return np.linalg.matrix_power(self.Phi, ell - 1) @ self.Gamma

    def sample_disturbance(self):
        """Sample a bounded disturbance for domain randomization."""
        d = np.random.uniform(-1, 1, size=3)
        d = d / (np.linalg.norm(d) + 1e-8) * self.d_bar * np.random.uniform()
        return d

    def update_params(self, n=None, dt=None, d_bar=None):
        if n is not None:
            self.n = n
            self.A_cw = self._build_A_cw(n)
        if dt is not None:
            self.dt = dt
        if d_bar is not None:
            self.d_bar = d_bar
        self.Phi, self.Gamma = self._build_exact_discrete(self.n, self.dt)


class NominalLQR:
    """
    Discrete LQR nominal controller: u_nom = -K (x_hat - x_h)
    Solves DARE for (Phi, Gamma, Q_lqr, R_lqr).  See Section 3.1.
    """

    def __init__(self, dynamics: CWDynamics, Q_lqr=None, R_lqr=None):
        self.dynamics = dynamics
        self.Q_lqr = Q_lqr if Q_lqr is not None else np.diag(
            [1.0, 1.0, 1.0, 0.1, 0.1, 0.1])
        self.R_lqr = R_lqr if R_lqr is not None else np.eye(3) * 10.0
        self.K = self._solve(dynamics.Phi, dynamics.Gamma,
                             self.Q_lqr, self.R_lqr)

    @staticmethod
    def _solve(Phi, Gamma, Q, R):
        P = solve_discrete_are(Phi, Gamma, Q, R)
        return np.linalg.solve(R + Gamma.T @ P @ Gamma, Gamma.T @ P @ Phi)

    def compute(self, x_hat, x_h):
        return -self.K @ (x_hat - x_h)

    def update(self, dynamics: CWDynamics):
        self.dynamics = dynamics
        self.K = self._solve(dynamics.Phi, dynamics.Gamma,
                             self.Q_lqr, self.R_lqr)


class KalmanFilter:
    """
    Discrete Kalman filter: o_k = H x_k + nu_k.  See Section 2.3.
    """

    def __init__(self, dynamics: CWDynamics, H=None, Sigma_o=None, Q_proc=None):
        self.dynamics = dynamics
        self.H = H if H is not None else np.eye(6)
        self.Sigma_o = Sigma_o if Sigma_o is not None else np.eye(6) * 1.0
        self.Q_proc = Q_proc if Q_proc is not None else np.eye(6) * 0.01
        self.x_hat = np.zeros(6)
        self.P = np.eye(6) * 100.0

    def predict(self, u):
        Phi = self.dynamics.Phi
        Gamma = self.dynamics.Gamma
        self.x_hat = Phi @ self.x_hat + Gamma @ u
        self.P = Phi @ self.P @ Phi.T + self.Q_proc

    def update(self, o):
        H = self.H
        S = H @ self.P @ H.T + self.Sigma_o
        K_kal = self.P @ H.T @ np.linalg.inv(S)
        self.x_hat = self.x_hat + K_kal @ (o - H @ self.x_hat)
        self.P = (np.eye(6) - K_kal @ H) @ self.P

    def reset(self, x0, P0=None):
        self.x_hat = x0.copy()
        self.P = P0 if P0 is not None else np.eye(6) * 100.0