"""
Magnetic structure and spin handling for PySpinW.

Handles:
- Magnetic atom definitions
- Spin structures (commensurate and incommensurate)
- Magnetic supercells
- Spin transformations
"""

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


@dataclass
class MagneticAtom:
    """
    Magnetic atom definition.
    
    Attributes
    ----------
    label : str
        Atom label (e.g., 'Fe3+', 'Cu2+')
    position : np.ndarray
        Fractional coordinates in unit cell, shape (3,)
    spin : float
        Spin quantum number S
    g_tensor : np.ndarray
        g-tensor, shape (3, 3), default is 2*I
    color : str
        Color for visualization
    formfactor : str or callable
        Magnetic form factor identifier or function
    """
    label: str = 'M'
    position: np.ndarray = field(default_factory=lambda: np.zeros(3))
    spin: float = 1.0
    g_tensor: np.ndarray = field(default_factory=lambda: 2.0 * np.eye(3))
    color: str = 'blue'
    formfactor: str = None
    
    def __post_init__(self):
        self.position = np.asarray(self.position)
        self.g_tensor = np.asarray(self.g_tensor)
        if self.g_tensor.shape == ():
            self.g_tensor = float(self.g_tensor) * np.eye(3)
    
    @property
    def moment(self) -> float:
        """Magnetic moment μ = g*S (in Bohr magnetons)."""
        return np.trace(self.g_tensor) / 3 * self.spin


@dataclass
class MagneticStructure:
    """
    Magnetic structure definition.
    
    Supports commensurate and incommensurate structures using
    the rotating coordinate system formalism.
    
    Attributes
    ----------
    k : np.ndarray
        Magnetic propagation vector, shape (3,)
    n : np.ndarray
        Normal vector for spiral plane, shape (3,)
    S : np.ndarray
        Spin vectors for each magnetic atom, shape (n_atoms, 3)
    n_ext : tuple
        Magnetic supercell extension (multiples of crystallographic cell)
    mode : str
        Structure mode: 'direct', 'helical', 'fourier', 'random'
    """
    k: np.ndarray = field(default_factory=lambda: np.zeros(3))
    n: np.ndarray = field(default_factory=lambda: np.array([0, 0, 1]))
    S: np.ndarray = None
    n_ext: Tuple[int, int, int] = (1, 1, 1)
    mode: str = 'direct'
    
    def __post_init__(self):
        self.k = np.asarray(self.k)
        self.n = np.asarray(self.n)
        self.n = self.n / np.linalg.norm(self.n)
        if self.S is not None:
            self.S = np.asarray(self.S)
    
    @property
    def is_commensurate(self) -> bool:
        """Check if structure is commensurate (k is integer or simple fraction)."""
        k_frac = self.k * np.prod(self.n_ext)
        return np.allclose(k_frac, np.round(k_frac))
    
    @property
    def is_collinear(self) -> bool:
        """Check if all spins are collinear."""
        if self.S is None or len(self.S) < 2:
            return True
        
        # Check if all spins are parallel or antiparallel
        S_norm = self.S / (np.linalg.norm(self.S, axis=1, keepdims=True) + 1e-10)
        dots = np.abs(S_norm @ S_norm[0])
        return np.allclose(dots, 1.0)
    
    @property
    def n_magnetic_atoms(self) -> int:
        """Number of magnetic atoms."""
        return len(self.S) if self.S is not None else 0
    
    def get_spins_in_supercell(self, n_ext: Tuple[int, int, int] = None) -> np.ndarray:
        """
        Get spin configuration in magnetic supercell.
        
        For helical structures, applies the rotation by k*R for each cell.
        
        Parameters
        ----------
        n_ext : tuple
            Supercell dimensions
            
        Returns
        -------
        S_ext : np.ndarray
            Spin vectors in supercell, shape (n_cells * n_atoms, 3)
        """
        if n_ext is None:
            n_ext = self.n_ext
        
        if self.S is None:
            return None
        
        n_atoms = len(self.S)
        n_cells = np.prod(n_ext)
        
        if self.is_commensurate or np.allclose(self.k, 0):
            # Simple replication for commensurate structures
            return np.tile(self.S, (n_cells, 1))
        else:
            # Helical rotation
            S_ext = np.zeros((n_cells * n_atoms, 3))
            idx = 0
            for i in range(n_ext[0]):
                for j in range(n_ext[1]):
                    for k in range(n_ext[2]):
                        R = np.array([i, j, k])
                        phase = 2 * np.pi * np.dot(self.k, R)
                        rot = _rotation_matrix_around_axis(self.n, phase)
                        for a in range(n_atoms):
                            S_ext[idx] = rot @ self.S[a]
                            idx += 1
            return S_ext
    
    def normalize_spins(self, magnitudes: np.ndarray = None):
        """Normalize spin vectors to given magnitudes (default: unit length)."""
        if self.S is None:
            return
        
        norms = np.linalg.norm(self.S, axis=1, keepdims=True)
        norms[norms < 1e-10] = 1.0
        
        if magnitudes is None:
            self.S = self.S / norms
        else:
            magnitudes = np.asarray(magnitudes).reshape(-1, 1)
            self.S = self.S / norms * magnitudes
    
    def rotate_to_frame(self, frame: str = 'xyz') -> np.ndarray:
        """
        Transform spins to specified coordinate frame.
        
        Parameters
        ----------
        frame : str
            Target frame: 'xyz' (Cartesian), 'nep' (normal-e-perpendicular)
            
        Returns
        -------
        S_rotated : np.ndarray
            Rotated spin vectors
        """
        if frame == 'xyz':
            return self.S.copy()
        elif frame == 'nep':
            # Rotate to frame where n is along z
            z = self.n
            x = np.cross(np.array([0, 0, 1]), z)
            if np.linalg.norm(x) < 1e-10:
                x = np.cross(np.array([0, 1, 0]), z)
            x = x / np.linalg.norm(x)
            y = np.cross(z, x)
            
            R = np.column_stack([x, y, z])
            return self.S @ R
        else:
            raise ValueError(f"Unknown frame: {frame}")


def _rotation_matrix_around_axis(axis: np.ndarray, angle: float) -> np.ndarray:
    """
    Create rotation matrix for rotation around an axis by given angle.
    
    Uses Rodrigues' rotation formula.
    
    Parameters
    ----------
    axis : np.ndarray
        Rotation axis (will be normalized), shape (3,)
    angle : float
        Rotation angle in radians
        
    Returns
    -------
    R : np.ndarray
        Rotation matrix, shape (3, 3)
    """
    axis = axis / np.linalg.norm(axis)
    K = np.array([
        [0, -axis[2], axis[1]],
        [axis[2], 0, -axis[0]],
        [-axis[1], axis[0], 0]
    ])
    R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
    return R


def generate_ferromagnetic(n_atoms: int, direction: np.ndarray = None) -> MagneticStructure:
    """
    Generate a ferromagnetic structure.
    
    Parameters
    ----------
    n_atoms : int
        Number of magnetic atoms
    direction : np.ndarray
        Spin direction (default: [0, 0, 1])
        
    Returns
    -------
    magstr : MagneticStructure
        Ferromagnetic structure
    """
    if direction is None:
        direction = np.array([0, 0, 1])
    direction = direction / np.linalg.norm(direction)
    
    S = np.tile(direction, (n_atoms, 1))
    return MagneticStructure(k=np.zeros(3), S=S, mode='direct')


def generate_antiferromagnetic(n_atoms: int, k: np.ndarray, 
                                direction: np.ndarray = None) -> MagneticStructure:
    """
    Generate an antiferromagnetic structure.
    
    Parameters
    ----------
    n_atoms : int
        Number of magnetic atoms
    k : np.ndarray
        Propagation vector (e.g., [0.5, 0, 0] for Neel)
    direction : np.ndarray
        Spin direction (default: [0, 0, 1])
        
    Returns
    -------
    magstr : MagneticStructure
        Antiferromagnetic structure
    """
    if direction is None:
        direction = np.array([0, 0, 1])
    direction = direction / np.linalg.norm(direction)
    
    S = np.zeros((n_atoms, 3))
    for i in range(n_atoms):
        phase = np.exp(1j * 2 * np.pi * i * np.dot(k, np.array([1, 0, 0])))
        S[i] = direction * np.real(phase)
    
    return MagneticStructure(k=k, S=S, mode='direct')


def generate_helical(n_atoms: int, k: np.ndarray, n: np.ndarray = None) -> MagneticStructure:
    """
    Generate a helical (spiral) structure.
    
    Parameters
    ----------
    n_atoms : int
        Number of magnetic atoms
    k : np.ndarray
        Propagation vector
    n : np.ndarray
        Normal to spiral plane (default: [0, 0, 1])
        
    Returns
    -------
    magstr : MagneticStructure
        Helical structure
    """
    if n is None:
        n = np.array([0, 0, 1])
    n = n / np.linalg.norm(n)
    
    # Initial spin perpendicular to n
    e1 = np.cross(n, np.array([1, 0, 0]))
    if np.linalg.norm(e1) < 1e-10:
        e1 = np.cross(n, np.array([0, 1, 0]))
    e1 = e1 / np.linalg.norm(e1)
    
    S = np.tile(e1, (n_atoms, 1))
    
    return MagneticStructure(k=k, n=n, S=S, mode='helical')


def generate_120_structure(n_atoms: int = 3) -> MagneticStructure:
    """
    Generate a 120° structure (triangular antiferromagnet ground state).
    
    Parameters
    ----------
    n_atoms : int
        Number of atoms (should be 3 for basic structure)
        
    Returns
    -------
    magstr : MagneticStructure
        120° structure with spins in xy-plane
    """
    angles = np.array([0, 2*np.pi/3, 4*np.pi/3])
    S = np.zeros((n_atoms, 3))
    for i in range(min(n_atoms, 3)):
        S[i] = [np.cos(angles[i]), np.sin(angles[i]), 0]
    
    # Repeat pattern if more atoms
    for i in range(3, n_atoms):
        S[i] = S[i % 3]
    
    return MagneticStructure(k=np.zeros(3), S=S, mode='direct')
