"""
Electric Field Gradient Calculations
Snow Dust Project – Spatial Gradient Analysis
SCIENTIFICALLY CORRECTED VERSION
"""

import numpy as np
from dataclasses import dataclass


@dataclass
class GradientResult:
    gradient_vector: np.ndarray    # ∇E [V/m²]
    gradient_magnitude: float      # |∇E| [V/m²]
    divergence: float              # ∇·E [V/m²]
    curl: np.ndarray               # ∇×E [V/m²]
    laplacian: float               # ∇²E [V/m³]
    direction: np.ndarray          # Unit vector


class FieldGradientAnalyzer:

    def __init__(self):
        # Default tetrahedron (1m edges)
        self.vertices = np.array([
            [1, 0, -1/np.sqrt(2)],
            [-1, 0, -1/np.sqrt(2)],
            [0, 1, 1/np.sqrt(2)],
            [0, -1, 1/np.sqrt(2)]
        ]) * 0.5
        
        # Pre-compute geometry
        self._compute_geometry()

    def _compute_geometry(self):
        """Pre-compute tetrahedron geometric properties"""
        v = self.vertices
        
        # Vectors from vertex 0
        v01 = v[1] - v[0]
        v02 = v[2] - v[0]
        v03 = v[3] - v[0]
        
        # Volume
        self.volume = abs(np.dot(v01, np.cross(v02, v03))) / 6.0
        
        # Face areas
        self.face_areas = np.zeros(4)
        self.face_normals = np.zeros((4, 3))
        
        # Face opposite each vertex
        faces = [(1,2,3), (0,3,2), (0,1,3), (0,2,1)]
        
        for i, (a,b,c) in enumerate(faces):
            ab = v[b] - v[a]
            ac = v[c] - v[a]
            normal = np.cross(ab, ac)
            area = 0.5 * np.linalg.norm(normal)
            self.face_areas[i] = area
            self.face_normals[i] = normal / (2 * area)

    # ==========================================================
    # GRADIENT CALCULATION (SCALAR FIELD)
    # ==========================================================
    def scalar_gradient(self, phi: np.ndarray) -> GradientResult:
        """
        Calculate gradient of scalar potential φ
        φ: array of 4 potential values at vertices [V]
        """
        # Finite volume method for tetrahedron
        grad = np.zeros(3)
        
        for i in range(4):
            # Average φ on face opposite vertex i
            other_vertices = [j for j in range(4) if j != i]
            phi_face = np.mean(phi[other_vertices])
            
            grad += phi_face * self.face_areas[i] * self.face_normals[i]
        
        grad /= (3 * self.volume)
        
        # Magnitude
        mag = np.linalg.norm(grad)
        
        # Direction
        direction = grad / mag if mag > 0 else np.zeros(3)
        
        # For scalar field: divergence = 0, curl = 0, laplacian = 0
        return GradientResult(
            grad, mag, 0.0, np.zeros(3), 0.0, direction
        )

    # ==========================================================
    # GRADIENT CALCULATION (VECTOR FIELD)
    # ==========================================================
    def vector_gradient(self, E: np.ndarray) -> GradientResult:
        """
        Calculate gradient of vector field E
        E: array of shape (4,3) [V/m] at vertices
        """
        # Jacobian matrix
        J = np.zeros((3, 3))
        
        # Calculate gradient for each component
        for i in range(3):
            phi_component = E[:, i]
            grad_i = self.scalar_gradient(phi_component).gradient_vector
            J[i, :] = grad_i
        
        # Gradient magnitude (Frobenius norm)
        grad_mag = np.linalg.norm(J, 'fro')
        
        # Divergence (trace)
        div = np.trace(J)
        
        # Curl
        curl = np.array([
            J[2, 1] - J[1, 2],
            J[0, 2] - J[2, 0],
            J[1, 0] - J[0, 1]
        ])
        
        # Laplacian (approximate)
        laplacian = 0.0
        for i in range(3):
            # Second derivative approximation
            laplacian += self._laplacian_component(E[:, i])
        
        # Principal direction (eigenvector of largest eigenvalue)
        eigenvalues, eigenvectors = np.linalg.eig(J)
        idx = np.argmax(np.abs(eigenvalues))
        direction = eigenvectors[:, idx].real
        direction = direction / np.linalg.norm(direction) if np.linalg.norm(direction) > 0 else direction
        
        return GradientResult(
            J.flatten(), grad_mag, div, curl, laplacian, direction
        )

    def _laplacian_component(self, phi: np.ndarray) -> float:
        """Approximate Laplacian of scalar component"""
        # Center of tetrahedron
        center = np.mean(self.vertices, axis=0)
        phi_center = np.mean(phi)
        
        laplace = 0.0
        for i in range(4):
            r = np.linalg.norm(self.vertices[i] - center)
            if r > 0:
                laplace += (phi[i] - phi_center) / (r ** 2)
        
        return 6 * laplace / 4.0

    # ==========================================================
    # FIELD LINE ANALYSIS
    # ==========================================================
    def field_line_properties(self, E_center: np.ndarray, grad_result: GradientResult) -> dict:
        """
        Analyze field line properties at center
        """
        E_mag = np.linalg.norm(E_center)
        
        if E_mag < 1e-12:
            return {
                'curvature': 0.0,
                'torsion': 0.0,
                'convergent': False,
                'solenoidal': False
            }
        
        E_dir = E_center / E_mag
        
        # Curvature: κ = |(E·∇)E| / |E|²
        if grad_result.gradient_vector.size == 9:
            J = grad_result.gradient_vector.reshape(3, 3)
            E_dot_grad_E = J @ E_center
            curvature = np.linalg.norm(E_dot_grad_E) / (E_mag ** 2)
        else:
            curvature = 0.0
        
        # Torsion
        torsion = np.dot(E_center, grad_result.curl) / (E_mag ** 2)
        
        # Convergence/divergence
        convergent = grad_result.divergence < 0
        solenoidal = abs(grad_result.divergence) < 1e-6
        
        return {
            'curvature': curvature,
            'torsion': torsion,
            'convergent': convergent,
            'solenoidal': solenoidal,
            'field_magnitude': E_mag
        }


# ==============================================================
# TEST
# ==============================================================
if __name__ == "__main__":
    analyzer = FieldGradientAnalyzer()
    
    print("Testing corrected field gradient analyzer\n")
    
    # Test 1: Linear potential field
    # φ = 1000 + 200x + 150y - 100z
    true_grad = np.array([200.0, 150.0, -100.0])
    
    phi = 1000.0 + np.dot(analyzer.vertices, true_grad)
    result = analyzer.scalar_gradient(phi)
    
    print(f"Test 1 - Linear potential:")
    print(f"  True gradient: {true_grad}")
    print(f"  Calculated: {result.gradient_vector}")
    print(f"  Error: {np.linalg.norm(result.gradient_vector - true_grad):.2e} V/m")
    
    # Test 2: Vector field
    E_field = np.tile([-200, -150, 100], (4, 1))  # E = -∇φ
    vec_result = analyzer.vector_gradient(E_field)
    
    print(f"\nTest 2 - Vector field:")
    print(f"  Gradient magnitude: {vec_result.gradient_magnitude:.1f} V/m²")
    print(f"  Divergence: {vec_result.divergence:.2e} V/m²")
    print(f"  Curl: {vec_result.curl}")
    
    # Field line analysis
    E_center = -true_grad
    props = analyzer.field_line_properties(E_center, vec_result)
    
    print(f"\nField line properties:")
    print(f"  Curvature: {props['curvature']:.2e} 1/m")
    print(f"  Torsion: {props['torsion']:.2e} 1/m")
    print(f"  Convergent: {props['convergent']}")
    
    print("\nModel ready for Snow Dust gradient-based navigation.")
