"""
Scale-Consistent Ground Verification Theory.

Implements:
- Dimensionless reference dynamics (Section 4.1, Theorem 2)
- Heterogeneous time-space scale mapping (Section 4.2)
- Error propagation bounds (Section 4.3, Theorem 3)

Reference: Sections 4.1–4.3, Theorems 2–3.
"""

import numpy as np
from scipy.linalg import expm


class DimensionlessModel:
    """
    Dimensionless CW reference dynamics.

    Defines:
        tilde_r = r / L_c
        tilde_v = v / (n L_c)
        tilde_u = u / (n^2 L_c)
        tau = n t

    The dimensionless system matrix tilde_A is universal (no scale parameters).
    """

    def __init__(self):
        # Universal dimensionless system matrix (Eq. in Section 4.1)
        self.A_tilde = np.array([
            [0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 1],
            [3, 0, 0, 0, 2, 0],
            [0, 0, 0, -2, 0, 0],
            [0, 0, -1, 0, 0, 0],
        ], dtype=np.float64)

        self.B_tilde = np.zeros((6, 3))
        self.B_tilde[3:6, :] = np.eye(3)

    def step(self, x_tilde, u_tilde, dtau):
        """Dimensionless step via matrix exponential."""
        Phi_tilde = expm(self.A_tilde * dtau)
        # Γ_tilde via integral approximation of expm
        # For exact: integrate expm(A*s) ds from 0 to dtau, applied to B
        # Use the augmented matrix trick
        n = 6
        m = 3
        M = np.zeros((n + m, n + m))
        M[:n, :n] = self.A_tilde * dtau
        M[:n, n:n + m] = self.B_tilde * dtau
        eM = expm(M)
        Gamma_tilde = eM[:n, n:n + m]
        return Phi_tilde @ x_tilde + Gamma_tilde @ u_tilde

    def to_dimensionless(self, r, v, u, L_c, n):
        """Convert physical quantities to dimensionless."""
        r_tilde = r / L_c
        v_tilde = v / (n * L_c)
        u_tilde = u / (n ** 2 * L_c)
        return r_tilde, v_tilde, u_tilde

    def to_physical(self, r_tilde, v_tilde, u_tilde, L_c, n):
        """Convert dimensionless quantities to physical."""
        r = r_tilde * L_c
        v = v_tilde * n * L_c
        u = u_tilde * n ** 2 * L_c
        return r, v, u

    def dimensionless_safety(self, rho_safe, L_c, theta_los):
        """
        Dimensionless safety constraint parameters.

        Returns
        -------
        rho_safe_tilde : float
        theta_los : float (unchanged — angles are dimensionless)
        """
        return rho_safe / L_c, theta_los


class ScaleMapping:
    """
    Heterogeneous time-space scale mapping between orbit and ground platform.

    Section 4.2:
        r_g = alpha_L * r_s
        t_g = alpha_T * t_s
        v_g = (alpha_L / alpha_T) * v_s
        u_g = (alpha_L / alpha_T^2) * u_s

    Parameters
    ----------
    L_s, n_s : float — space task characteristic length and frequency
    L_g, n_g : float — ground platform characteristic length and frequency
    """

    def __init__(self, L_s, n_s, L_g, n_g):
        self.L_s = L_s
        self.n_s = n_s
        self.L_g = L_g
        self.n_g = n_g

        self.alpha_L = L_g / L_s
        self.alpha_T = n_s / n_g

    def space_to_ground(self, r_s, v_s, u_s):
        """Map space-task quantities to ground platform."""
        r_g = self.alpha_L * r_s
        v_g = (self.alpha_L / self.alpha_T) * v_s
        u_g = (self.alpha_L / self.alpha_T ** 2) * u_s
        return r_g, v_g, u_g

    def ground_to_space(self, r_g, v_g, u_g):
        """Map ground platform quantities to space task."""
        r_s = r_g / self.alpha_L
        v_s = v_g * self.alpha_T / self.alpha_L
        u_s = u_g * self.alpha_T ** 2 / self.alpha_L
        return r_s, v_s, u_s

    def check_feasibility(self, v_s_max, u_s_max, v_g_max, u_g_max):
        """
        Check if scale mapping is feasible for ground platform capabilities.

        Returns
        -------
        feasible : bool
        info : dict with margin details
        """
        v_g_req = (self.alpha_L / self.alpha_T) * v_s_max
        u_g_req = (self.alpha_L / self.alpha_T ** 2) * u_s_max

        feasible = (v_g_req <= v_g_max) and (u_g_req <= u_g_max)
        info = {
            'v_g_required': v_g_req,
            'u_g_required': u_g_req,
            'v_margin': v_g_max - v_g_req,
            'u_margin': u_g_max - u_g_req,
            'feasible': feasible,
        }
        return feasible, info


class ErrorPropagation:
    """
    Finite-horizon error propagation bound (Theorem 3).

    ||delta(tau)|| <= C_H * tau_H * sup ||e_a_tilde(sigma)||

    where C_H = sup_{0 <= s <= tau_H} ||exp(A_tilde s) B_tilde||_2
    """

    def __init__(self, tau_H=100.0, n_samples=500):
        self.model = DimensionlessModel()
        self.tau_H = tau_H
        self.C_H = self._compute_C_H(tau_H, n_samples)

    def _compute_C_H(self, tau_H, n_samples):
        """Compute C_H by sampling the matrix exponential."""
        A = self.model.A_tilde
        B = self.model.B_tilde
        C_H = 0.0
        for s in np.linspace(0, tau_H, n_samples):
            eAs = expm(A * s)
            val = np.linalg.norm(eAs @ B, 2)
            C_H = max(C_H, val)
        return C_H

    def bound(self, e_a_sup):
        """
        Compute the error bound.

        Parameters
        ----------
        e_a_sup : float
            Supremum of dimensionless acceleration tracking error.

        Returns
        -------
        delta_bound : float
            Upper bound on ||delta(tau)|| for tau in [0, tau_H].
        """
        return self.C_H * self.tau_H * e_a_sup

    def convert_physical_error(self, e_a_phys, n_g, L_g):
        """
        Convert physical acceleration error to dimensionless.

        e_a_tilde = e_a / (n_g^2 L_g)
        """
        return e_a_phys / (n_g ** 2 * L_g)