"""
Linear Spin Wave Theory (LSWT) module.

Improved implementation based on spinw_torch minimalist port.
Provides cleaner Holstein-Primakoff linearization and efficient diagonalization.
"""

import numpy as np
from typing import Tuple, Optional, Dict, List
from dataclasses import dataclass


@dataclass
class LSWTResult:
    """Container for LSWT calculation results."""
    omega: np.ndarray      # (n_q, n_modes) energies
    V: np.ndarray          # (n_q, 2n, 2n) transformation matrices
    Sab: np.ndarray        # (n_q, 3, 3, n_modes) correlation tensor
    hkl: np.ndarray        # (n_q, 3) Q-points


def build_rotation_matrices(directions: np.ndarray) -> np.ndarray:
    """
    Build local-to-global rotation matrices R_i.
    
    Parameters
    ----------
    directions : ndarray, shape (n_atoms, 3)
        Classical spin directions (will be normalized)
        
    Returns
    -------
    R : ndarray, shape (n_atoms, 3, 3)
        Rotation matrices where columns are (e1, e2, e3) with e3 = spin direction
    """
    n = directions.shape[0]
    
    # Normalize directions to get e3 (z-axis in local frame)
    ez = directions / (np.linalg.norm(directions, axis=-1, keepdims=True) + 1e-10)
    
    # Reference vectors for building orthogonal basis
    ref1 = np.array([1.0, 0.0, 0.0])
    ref2 = np.array([0.0, 1.0, 0.0])
    
    # Choose reference based on alignment with ez
    dot1 = np.abs(np.sum(ez * ref1, axis=-1))
    ref = np.where(dot1[:, None] > 0.9, ref2[None, :], ref1[None, :])
    
    # Build orthonormal basis
    e1 = np.cross(ez, ref)
    e1 = e1 / (np.linalg.norm(e1, axis=-1, keepdims=True) + 1e-10)
    e2 = np.cross(ez, e1)
    e2 = e2 / (np.linalg.norm(e2, axis=-1, keepdims=True) + 1e-10)
    
    # Stack as rotation matrix [e1 | e2 | ez]
    R = np.stack([e1, e2, ez], axis=-1)  # (n, 3, 3)
    
    return R


def build_hp_matrices(spins: np.ndarray) -> np.ndarray:
    """
    Build Holstein-Primakoff linearization matrices.
    
    For transverse spin operators:
        S_x = sqrt(S/2) * (a + a†)
        S_y = sqrt(S/2) * (-i*a + i*a†)
    
    Parameters
    ----------
    spins : ndarray, shape (n_atoms,)
        Spin quantum numbers S_i
        
    Returns
    -------
    M : ndarray, shape (n_atoms, 2, 2), complex
        M[i, alpha, op] where alpha=0,1 for x,y and op=0,1 for a, a†
    """
    n = len(spins)
    sqrtS = np.sqrt(spins / 2.0)
    
    M = np.zeros((n, 2, 2), dtype=np.complex128)
    M[:, 0, 0] = sqrtS        # S_x, a
    M[:, 0, 1] = sqrtS        # S_x, a†
    M[:, 1, 0] = -1j * sqrtS  # S_y, a  
    M[:, 1, 1] = 1j * sqrtS   # S_y, a†
    
    return M


def spinwave_lswt(
    positions: np.ndarray,
    spins: np.ndarray,
    directions: np.ndarray,
    couplings: List[Dict],
    q_points: np.ndarray,
    k_mag: np.ndarray = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute LSWT magnon energies using the efficient σz @ H method.
    
    This implements bosonic Bogoliubov diagonalization following the
    approach in the spinw_torch minimalist port.
    
    Parameters
    ----------
    positions : ndarray, shape (n_atoms, 3)
        Fractional coordinates of magnetic atoms
    spins : ndarray, shape (n_atoms,)
        Spin quantum numbers S_i
    directions : ndarray, shape (n_atoms, 3)
        Classical spin directions
    couplings : list of dict
        Each coupling has keys: 'i', 'j', 'delta', 'J'
        - i, j: atom indices (0-based)
        - delta: lattice translation (3,)
        - J: exchange matrix (3, 3) or scalar
    q_points : ndarray, shape (n_q, 3)
        Q-points in reciprocal lattice units
    k_mag : ndarray, optional
        Magnetic propagation vector (default: [0,0,0])
        
    Returns
    -------
    omega : ndarray, shape (n_q, n_atoms)
        Magnon energies (positive branches)
    eigenvectors : ndarray, shape (n_q, 2n, 2n)
        Transformation matrices
    """
    if k_mag is None:
        k_mag = np.array([0.0, 0.0, 0.0])
    
    n_atoms = len(spins)
    n_q = len(q_points)
    
    # Build rotation matrices and HP matrices
    R = build_rotation_matrices(directions)
    M = build_hp_matrices(spins)
    
    # Allocate A(q) and B(q) matrices
    A = np.zeros((n_q, n_atoms, n_atoms), dtype=np.complex128)
    B = np.zeros((n_q, n_atoms, n_atoms), dtype=np.complex128)
    
    two_pi = 2.0 * np.pi
    
    # Remove duplicate bonds: keep only (i,j,delta) where delta >= 0 lexicographically
    # or i < j for intra-cell bonds
    unique_couplings = []
    seen = set()
    
    for coup in couplings:
        i, j = coup['i'], coup['j']
        delta = tuple(coup['delta'])
        
        # Create canonical key: sort by (delta, i, j) to identify duplicates
        # Bond (i, j, delta) is equivalent to (j, i, -delta)
        neg_delta = tuple(-x for x in delta)
        
        key1 = (i, j, delta)
        key2 = (j, i, neg_delta)
        
        # Use lexicographic ordering to pick canonical form
        if key1 not in seen and key2 not in seen:
            seen.add(key1)
            seen.add(key2)
            unique_couplings.append(coup)
    
    # Process each unique coupling
    for coup in unique_couplings:
        i = coup['i']
        j = coup['j']
        delta = np.array(coup['delta'], dtype=np.float64)
        
        # Get exchange matrix (scalar -> isotropic)
        J = coup['J']
        if np.isscalar(J):
            J = float(J) * np.eye(3)
        else:
            J = np.asarray(J, dtype=np.float64)
        
        # Bond vector in fractional coordinates
        d = delta + positions[j] - positions[i]
        
        # Phase factor: exp(2πi * (q - k_mag) · d)
        q_eff = q_points - k_mag[None, :]
        qdot = np.sum(q_eff * d[None, :], axis=-1)
        phase = np.exp(1j * two_pi * qdot)
        phase_conj = np.conj(phase)
        
        # Transform exchange to local frames: J_loc = R_i^T @ J @ R_j
        J_local = R[i].T @ J @ R[j]
        
        # Extract blocks
        Jxy = J_local[:2, :2]  # Transverse part (2x2)
        Jzz = J_local[2, 2]    # Longitudinal part (scalar)
        
        # Coefficient matrix C = M_i^T @ Jxy @ M_j
        # NOTE: Use regular transpose, NOT conjugate transpose!
        # This encodes the structure of bilinear operators:
        #   C[0,0] = coeff of a_i a_j
        #   C[0,1] = coeff of a_i a†_j  
        #   C[1,0] = coeff of a†_i a_j
        #   C[1,1] = coeff of a†_i a†_j
        C = M[i].T @ Jxy @ M[j]  # (2, 2)
        
        c_dag_a = C[1, 0]     # a_i† a_j coefficient
        c_a_dag = C[0, 1]     # a_i a_j† coefficient  
        c_dag_dag = C[1, 1]   # a_i† a_j† coefficient
        
        # Normal (A) terms: a†_i a_j type
        # Include both (i,j) and (j,i) contributions from the bond
        A[:, i, j] += c_dag_a * phase
        A[:, j, i] += c_a_dag * phase_conj
        
        # Anomalous (B) terms: a†_i a†_j type
        B[:, i, j] += c_dag_dag * phase
        B[:, j, i] += c_dag_dag * phase_conj
        
        # Longitudinal contribution: -Jzz * S_j * n_i - Jzz * S_i * n_j
        # (from Sz_i Sz_j term in Holstein-Primakoff)
        A[:, i, i] -= Jzz * spins[j]
        A[:, j, j] -= Jzz * spins[i]
    
    # Build full BdG Hamiltonian H(q) of shape (n_q, 2n, 2n)
    # H = [[A, B], [B*, A*]]
    n2 = 2 * n_atoms
    H = np.zeros((n_q, n2, n2), dtype=np.complex128)
    H[:, :n_atoms, :n_atoms] = A
    H[:, :n_atoms, n_atoms:] = B
    H[:, n_atoms:, :n_atoms] = np.conj(B).transpose(0, 2, 1)
    H[:, n_atoms:, n_atoms:] = np.conj(A)
    
    # Diagonalize using σz @ H method
    # M = σz @ H where σz = diag(I, -I)
    M_diag = H.copy()
    M_diag[:, n_atoms:, :] = -M_diag[:, n_atoms:, :]
    
    # Eigenvalue decomposition
    eigenvalues, eigenvectors = np.linalg.eig(M_diag)
    
    # Sort by real part (positive eigenvalues are magnon energies)
    idx = np.argsort(np.real(eigenvalues), axis=-1)[:, ::-1]
    
    omega_all = np.zeros((n_q, n_atoms), dtype=np.float64)
    V_all = np.zeros((n_q, n2, n2), dtype=np.complex128)
    
    for iq in range(n_q):
        sorted_idx = idx[iq]
        omega_all[iq] = np.real(eigenvalues[iq, sorted_idx[:n_atoms]])
        V_all[iq] = eigenvectors[iq, :, sorted_idx]
    
    # Clamp small negative values (numerical noise)
    omega_all = np.maximum(omega_all, 0.0)
    
    return omega_all, V_all


def calculate_correlation_tensor(
    V: np.ndarray,
    spins: np.ndarray,
    R: np.ndarray,
    q_points: np.ndarray,
) -> np.ndarray:
    """
    Calculate the spin-spin correlation tensor S^{αβ}(q, ω).
    
    Parameters
    ----------
    V : ndarray, shape (n_q, 2n, 2n)
        Transformation matrices from diagonalization
    spins : ndarray, shape (n_atoms,)
        Spin quantum numbers
    R : ndarray, shape (n_atoms, 3, 3)
        Local-to-global rotation matrices
    q_points : ndarray, shape (n_q, 3)
        Q-points
        
    Returns
    -------
    Sab : ndarray, shape (n_q, 3, 3, n_atoms)
        Correlation tensor
    """
    n_q = len(q_points)
    n_atoms = len(spins)
    
    # Placeholder - full implementation requires careful treatment
    # of the Bogoliubov transformation
    Sab = np.zeros((n_q, 3, 3, n_atoms), dtype=np.complex128)
    
    # Simple approximation: transverse correlations scale with sqrt(S)
    for i in range(n_atoms):
        sqrtS = np.sqrt(spins[i] / 2)
        
        # Get local axes
        e1 = R[i, :, 0]
        e2 = R[i, :, 1]
        
        # Transverse correlations
        for a in range(3):
            for b in range(3):
                Sab[:, a, b, i] = sqrtS * (e1[a] * e1[b] + e2[a] * e2[b])
    
    return Sab


def q_path(q_start, q_end, n_pts: int) -> np.ndarray:
    """
    Generate a straight-line q-path between two points.
    
    Parameters
    ----------
    q_start, q_end : array-like
        Start and end points in reciprocal lattice units
    n_pts : int
        Number of points (inclusive)
        
    Returns
    -------
    q : ndarray, shape (n_pts, 3)
        Q-points along the path
    """
    q_start = np.asarray(q_start)
    q_end = np.asarray(q_end)
    
    t = np.linspace(0.0, 1.0, n_pts)
    q = q_start[None, :] * (1.0 - t[:, None]) + q_end[None, :] * t[:, None]
    
    return q


def multi_segment_path(waypoints: List, n_pts: int) -> np.ndarray:
    """
    Generate a multi-segment Q-path through waypoints.
    
    Parameters
    ----------
    waypoints : list of array-like
        List of Q-points to connect
    n_pts : int
        Total number of points
        
    Returns
    -------
    q : ndarray, shape (n_pts, 3)
        Q-points along the path
    """
    n_segments = len(waypoints) - 1
    pts_per_segment = n_pts // n_segments
    
    paths = []
    for i in range(n_segments):
        endpoint = (i == n_segments - 1)
        segment = q_path(waypoints[i], waypoints[i + 1], 
                        pts_per_segment + (1 if endpoint else 0))
        if not endpoint:
            segment = segment[:-1]  # Remove endpoint to avoid duplication
        paths.append(segment)
    
    return np.vstack(paths)
