"""
Main SpinW class for spin wave calculations.

This is the central class that combines:
- Crystal lattice definition
- Magnetic atoms and structure
- Exchange couplings
- Linear spin wave theory calculations
- Neutron scattering cross-sections
"""

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

from .lattice import Lattice, SpaceGroup
from .magnetic import MagneticAtom, MagneticStructure
from .coupling import Coupling, ExchangeMatrix, Bond
from .constants import CONSTANTS
from ..backends import current_backend


class SpinW:
    """
    Main class for spin wave calculations.
    
    SpinW (spin-double-u) is a Python library for calculating spin wave
    dispersion and neutron scattering cross-sections using linear spin
    wave theory.
    
    Examples
    --------
    >>> # Create a triangular lattice antiferromagnet
    >>> sw = SpinW()
    >>> sw.genlattice(lat_const=[3, 3, 6], angled=[90, 90, 120])
    >>> sw.addatom(r=[0, 0, 0], S=1, label='Cu1')
    >>> sw.gencoupling(max_distance=5)
    >>> sw.addmatrix(label='J1', value=1.0)
    >>> sw.addcoupling(mat='J1', bond=1)
    >>> sw.genmagstr(mode='helical', k=[1/3, 1/3, 0])
    >>> spec = sw.spinwave([[0, 0, 0], [1, 1, 0]], n_pts=100)
    """
    
    def __init__(self):
        """Initialize empty SpinW object."""
        self._lattice = Lattice()
        self._atoms = []  # List of MagneticAtom
        self._magstr = MagneticStructure()
        self._coupling = Coupling()
        
        # Backend for computations
        self._backend = current_backend()
        
        # Cached data
        self._cache = {}
    
    def copy(self) -> 'SpinW':
        """Create a deep copy of this SpinW object."""
        return copy.deepcopy(self)
    
    # ========== Lattice methods ==========
    
    def genlattice(self, lat_const: List[float] = None, 
                   angled: List[float] = None,
                   spgr: Union[int, str] = 'P 1',
                   sym: Union[int, str] = None,
                   **kwargs) -> 'SpinW':
        """
        Generate crystal lattice.
        
        Parameters
        ----------
        lat_const : list
            Lattice constants [a, b, c] in Angstrom
        angled : list
            Lattice angles [alpha, beta, gamma] in degrees
        spgr : int or str
            Space group number or symbol
        sym : int or str
            Alias for spgr
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        if lat_const is None:
            lat_const = [1.0, 1.0, 1.0]
        if angled is None:
            angled = [90.0, 90.0, 90.0]
        
        # Handle sym alias
        if sym is not None:
            spgr = sym
        
        # Create space group
        if isinstance(spgr, int):
            if spgr == 0:
                sg = SpaceGroup()  # P1, no symmetry
            else:
                sg = SpaceGroup.from_number(spgr)
        else:
            sg = SpaceGroup.from_symbol(str(spgr))
        
        self._lattice = Lattice(
            lat_const=tuple(lat_const),
            angles=tuple(angled),
            spacegroup=sg
        )
        
        self._clear_cache()
        return self
    
    @property
    def lattice(self) -> Lattice:
        """Get the lattice object."""
        return self._lattice
    
    # ========== Atom methods ==========
    
    def addatom(self, r: List[float], S: float = 0, 
                label: str = None, color: str = 'blue',
                g: Union[float, np.ndarray] = 2.0,
                formfactn: Union[str, float] = None,
                **kwargs) -> 'SpinW':
        """
        Add a magnetic atom to the crystal structure.
        
        Parameters
        ----------
        r : list
            Fractional coordinates [x, y, z]
        S : float
            Spin quantum number (0 for non-magnetic)
        label : str
            Atom label (e.g., 'Fe3+', 'Cu2+')
        color : str
            Color for visualization
        g : float or array
            g-factor or g-tensor (default: 2)
        formfactn : str or float
            Magnetic form factor
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        r = np.asarray(r)
        
        # Handle multiple positions
        if r.ndim == 2:
            for i in range(r.shape[1] if r.shape[0] == 3 else r.shape[0]):
                pos = r[:, i] if r.shape[0] == 3 else r[i]
                self.addatom(r=pos, S=S, label=label, color=color, g=g)
            return self
        
        if label is None:
            label = f'atom{len(self._atoms) + 1}'
        
        # Handle g-tensor
        if np.isscalar(g):
            g_tensor = float(g) * np.eye(3)
        else:
            g_tensor = np.asarray(g)
            if g_tensor.shape == (3,):
                g_tensor = np.diag(g_tensor)
        
        atom = MagneticAtom(
            label=label,
            position=r,
            spin=S,
            g_tensor=g_tensor,
            color=color,
            formfactor=formfactn
        )
        
        self._atoms.append(atom)
        self._clear_cache()
        return self
    
    def atom(self, idx: int = None) -> Union[List[MagneticAtom], MagneticAtom]:
        """Get atom(s) by index."""
        if idx is None:
            return self._atoms
        return self._atoms[idx]
    
    def matom(self, idx: int = None) -> Union[List[MagneticAtom], MagneticAtom]:
        """Get magnetic atom(s) by index."""
        mag_atoms = [a for a in self._atoms if a.spin > 0]
        if idx is None:
            return mag_atoms
        return mag_atoms[idx]
    
    @property
    def n_atoms(self) -> int:
        """Total number of atoms."""
        return len(self._atoms)
    
    @property
    def n_mag_atoms(self) -> int:
        """Number of magnetic atoms."""
        return len([a for a in self._atoms if a.spin > 0])
    
    # ========== Coupling methods ==========
    
    def gencoupling(self, max_distance: float, **kwargs) -> 'SpinW':
        """
        Generate magnetic bonds up to a maximum distance.
        
        Parameters
        ----------
        max_distance : float
            Maximum bond distance in Angstrom
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        positions = np.array([a.position for a in self._atoms])
        self._coupling.generate_bonds(positions, self._lattice, max_distance)
        return self
    
    def addmatrix(self, label: str, value: Union[float, np.ndarray],
                  color: str = 'gray', **kwargs) -> 'SpinW':
        """
        Add an exchange matrix.
        
        Parameters
        ----------
        label : str
            Matrix label (e.g., 'J1', 'J2', 'A')
        value : float or array
            Exchange value or matrix
        color : str
            Visualization color
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        self._coupling.add_matrix(label, value, color)
        return self
    
    def addcoupling(self, mat: str, bond: int, **kwargs) -> 'SpinW':
        """
        Assign exchange matrix to a bond type.
        
        Parameters
        ----------
        mat : str
            Exchange matrix label
        bond : int
            Bond type index
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        self._coupling.add_coupling(mat, bond)
        return self
    
    def addaniso(self, mat: str, atom_idx: int = None, **kwargs) -> 'SpinW':
        """
        Add single-ion anisotropy.
        
        Parameters
        ----------
        mat : str
            Anisotropy matrix label
        atom_idx : int
            Atom index (or all magnetic atoms if None)
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        if atom_idx is None:
            for i, a in enumerate(self._atoms):
                if a.spin > 0:
                    self._coupling.add_anisotropy(mat, i)
        else:
            self._coupling.add_anisotropy(mat, atom_idx)
        return self
    
    def field(self, B: List[float]) -> 'SpinW':
        """
        Set external magnetic field.
        
        Parameters
        ----------
        B : list
            Magnetic field [Bx, By, Bz] in Tesla
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        self._coupling.set_field(np.asarray(B))
        return self
    
    def temperature(self, T: float) -> 'SpinW':
        """
        Set temperature.
        
        Parameters
        ----------
        T : float
            Temperature in Kelvin
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        self._coupling.set_temperature(T)
        return self
    
    def table(self, what: str = 'atom') -> str:
        """
        Print information table.
        
        Parameters
        ----------
        what : str
            'atom', 'matom', 'bond', 'matrix'
        """
        if what == 'atom':
            lines = ["Atoms:"]
            lines.append(f"{'idx':>4} {'label':>8} {'S':>5} {'position':>20}")
            for i, a in enumerate(self._atoms):
                pos_str = f"[{a.position[0]:.3f}, {a.position[1]:.3f}, {a.position[2]:.3f}]"
                lines.append(f"{i:4d} {a.label:>8} {a.spin:5.1f} {pos_str:>20}")
            return "\n".join(lines)
        
        elif what == 'matom':
            mag_atoms = [a for a in self._atoms if a.spin > 0]
            lines = ["Magnetic atoms:"]
            for i, a in enumerate(mag_atoms):
                lines.append(f"  {i}: {a.label}, S={a.spin}")
            return "\n".join(lines)
        
        elif what == 'bond':
            return self._coupling.list_bonds()
        
        elif what == 'matrix':
            lines = ["Exchange matrices:"]
            for label, mat in self._coupling.exchange_matrices.items():
                lines.append(f"  {label}: {mat.isotropic_part:.4f} (isotropic)")
            return "\n".join(lines)
        
        return ""
    
    # ========== Magnetic structure methods ==========
    
    def genmagstr(self, mode: str = 'direct', k: List[float] = None,
                  n: List[float] = None, S: np.ndarray = None,
                  n_ext: Tuple[int, int, int] = None,
                  unit: str = 'lu', **kwargs) -> 'SpinW':
        """
        Generate magnetic structure.
        
        Parameters
        ----------
        mode : str
            'direct': specify all spins directly
            'helical': helical/spiral structure
            'random': random spin configuration
            'fourier': Fourier components
        k : list
            Propagation vector [h, k, l]
        n : list
            Normal vector for spiral plane
        S : array
            Spin vectors, shape (n_atoms, 3) or (3, n_atoms)
        n_ext : tuple
            Supercell dimensions
        unit : str
            'lu' for lattice units, 'xyz' for Cartesian
            
        Returns
        -------
        self : SpinW
            For method chaining
        """
        mag_atoms = [a for a in self._atoms if a.spin > 0]
        n_mag = len(mag_atoms)
        
        if k is None:
            k = [0, 0, 0]
        k = np.asarray(k)
        
        if n is None:
            n = [0, 0, 1]
        n = np.asarray(n)
        n = n / np.linalg.norm(n)
        
        if n_ext is None:
            n_ext = (1, 1, 1)
        
        if mode == 'direct':
            if S is None:
                # Default to z-polarized
                S = np.zeros((n_mag, 3))
                S[:, 2] = 1.0
            else:
                S = np.asarray(S)
                if S.shape[0] == 3 and S.shape[1] != 3:
                    S = S.T
        
        elif mode == 'helical':
            # Initialize spins perpendicular to n
            e1 = np.cross(n, [1, 0, 0])
            if np.linalg.norm(e1) < 1e-10:
                e1 = np.cross(n, [0, 1, 0])
            e1 = e1 / np.linalg.norm(e1)
            
            if S is None:
                S = np.tile(e1, (n_mag, 1))
            else:
                S = np.asarray(S)
                if S.shape[0] == 3 and S.shape[1] != 3:
                    S = S.T
        
        elif mode == 'random':
            # Random spin directions
            S = np.random.randn(n_mag, 3)
            norms = np.linalg.norm(S, axis=1, keepdims=True)
            S = S / norms
        
        elif mode == 'fourier':
            raise NotImplementedError("Fourier mode not yet implemented")
        
        # Normalize spins by spin quantum numbers
        spin_mags = np.array([a.spin for a in mag_atoms])
        S = S / (np.linalg.norm(S, axis=1, keepdims=True) + 1e-10)
        S = S * spin_mags[:, np.newaxis]
        
        self._magstr = MagneticStructure(k=k, n=n, S=S, n_ext=n_ext, mode=mode)
        self._clear_cache()
        return self
    
    @property
    def magstr(self) -> MagneticStructure:
        """Get the magnetic structure."""
        return self._magstr
    
    # ========== Energy and optimization ==========
    
    def energy(self) -> float:
        """
        Calculate the classical energy of the current spin configuration.
        
        Returns
        -------
        E : float
            Energy per spin in meV
        """
        from .coupling import quick_hamiltonian
        E = quick_hamiltonian(self._magstr.S, self._coupling, self._lattice, 
                              k=self._magstr.k, n=self._magstr.n)
        return E / self.n_mag_atoms
    
    # ========== Spin wave calculation ==========
    
    def spinwave(self, Q: Union[List, np.ndarray], 
                 n_pts: int = 100,
                 hermit: bool = True,
                 formfact: bool = False,
                 **kwargs) -> Dict:
        """
        Calculate spin wave dispersion.
        
        This is the main spin wave calculation using linear spin wave theory.
        Diagonalizes the spin wave Hamiltonian at each Q-point.
        
        Parameters
        ----------
        Q : list or array
            Q-points specification:
            - List of [start, end] for linear scan
            - List of [[h1,k1,l1], [h2,k2,l2], n_pts] for linear scan
            - Array of shape (n_q, 3) for explicit Q-points
        n_pts : int
            Number of points for linear scan
        hermit : bool
            If True, assume Hermitian Hamiltonian (faster)
        formfact : bool
            Include magnetic form factor
            
        Returns
        -------
        spec : dict
            Dictionary containing:
            - 'omega': Spin wave energies, shape (n_q, n_modes)
            - 'Sab': Spin-spin correlation tensor, shape (n_q, 3, 3, n_modes)
            - 'hkl': Q-points, shape (n_q, 3)
            - 'hklA': Q in Angstrom^-1, shape (n_q, 3)
        """
        # Parse Q specification
        hkl = self._parse_q_path(Q, n_pts)
        n_q = len(hkl)
        
        # Get magnetic atoms
        mag_atoms = [a for a in self._atoms if a.spin > 0]
        n_mag = len(mag_atoms)
        
        if n_mag == 0:
            raise ValueError("No magnetic atoms defined")
        
        if self._magstr.S is None:
            raise ValueError("Magnetic structure not defined. Call genmagstr() first.")
        
        # Build and diagonalize Hamiltonian
        omega, V, Sab = self._calculate_spinwave(hkl, hermit=hermit)
        
        # Calculate Q in Angstrom^-1
        hklA = np.array([self._lattice.hkl_to_q(q) for q in hkl])
        
        # Apply form factor if requested
        if formfact:
            ff = self._calculate_formfactor(hklA)
            Sab = Sab * ff[:, np.newaxis, np.newaxis, np.newaxis]
        
        return {
            'omega': omega,
            'Sab': Sab,
            'V': V,
            'hkl': hkl,
            'hklA': hklA,
            'obj': self,
        }
    
    def _parse_q_path(self, Q, n_pts: int) -> np.ndarray:
        """Parse Q-point specification into array of hkl values."""
        if isinstance(Q, np.ndarray) and Q.ndim == 2:
            return Q
        
        if isinstance(Q, list):
            # Check if all elements are 3-vectors
            all_3vec = all(len(q) == 3 for q in Q if hasattr(q, '__len__'))
            
            if len(Q) == 2 and all_3vec:
                # [start, end] format
                start = np.array(Q[0])
                end = np.array(Q[1])
                return np.linspace(start, end, n_pts)
            
            if len(Q) >= 3 and all_3vec:
                # Multiple waypoints - multi-segment path
                paths = []
                n_segments = len(Q) - 1
                pts_per_segment = max(2, n_pts // n_segments)
                
                for i in range(n_segments):
                    start = np.array(Q[i])
                    end = np.array(Q[i + 1])
                    endpoint = (i == n_segments - 1)
                    segment = np.linspace(start, end, pts_per_segment + (1 if endpoint else 0))
                    if not endpoint:
                        segment = segment[:-1]
                    paths.append(segment)
                return np.vstack(paths)
            
            if len(Q) == 3 and np.isscalar(Q[2]):
                # [start, end, n_pts] format (deprecated)
                start = np.array(Q[0])
                end = np.array(Q[1])
                return np.linspace(start, end, int(Q[2]))
        
        raise ValueError(f"Invalid Q specification: {Q}")
    
    def _calculate_spinwave(self, hkl: np.ndarray, hermit: bool = True) -> Tuple:
        """
        Core spin wave calculation using improved LSWT.
        
        Uses the efficient σz @ H diagonalization approach from spinw_torch.
        """
        from .lswt import spinwave_lswt, build_rotation_matrices, calculate_correlation_tensor
        
        backend = self._backend
        n_q = len(hkl)
        
        mag_atoms = [a for a in self._atoms if a.spin > 0]
        n_mag = len(mag_atoms)
        
        # Extract positions, spins, and directions
        positions = np.array([a.position for a in mag_atoms])
        spins = np.array([a.spin for a in mag_atoms])
        
        # Get spin directions from magnetic structure
        S = self._magstr.S
        directions = S / (np.linalg.norm(S, axis=1, keepdims=True) + 1e-10)
        
        # Build coupling list in the format expected by LSWT
        couplings = []
        for bond in self._coupling.bonds:
            if bond.exchange is None:
                continue
            
            # Map to magnetic atom indices
            mag_idx_i = self._get_mag_atom_index(bond.atom1)
            mag_idx_j = self._get_mag_atom_index(bond.atom2)
            
            if mag_idx_i is None or mag_idx_j is None:
                continue
            
            couplings.append({
                'i': mag_idx_i,
                'j': mag_idx_j,
                'delta': tuple(bond.dl),
                'J': bond.exchange.matrix
            })
        
        # Call the improved LSWT calculation
        omega, V = spinwave_lswt(
            positions=positions,
            spins=spins,
            directions=directions,
            couplings=couplings,
            q_points=hkl,
            k_mag=self._magstr.k
        )
        
        # Calculate correlation tensor
        R = build_rotation_matrices(directions)
        Sab = calculate_correlation_tensor(V, spins, R, hkl)
        
        return omega, V, Sab
    
    def _build_hamiltonian_blocks(self, q: np.ndarray, S: np.ndarray, 
                                   mag_atoms: List[MagneticAtom]) -> Tuple:
        """Build A and B matrices for the spin wave Hamiltonian."""
        n_mag = len(mag_atoms)
        
        # Initialize
        A = np.zeros((n_mag, n_mag), dtype=np.complex128)
        B = np.zeros((n_mag, n_mag), dtype=np.complex128)
        
        # Exchange contributions
        for bond in self._coupling.bonds:
            if bond.exchange is None:
                continue
            
            # Check if both atoms are magnetic
            i = bond.atom1
            j = bond.atom2
            
            # Map to magnetic atom indices
            mag_idx_i = self._get_mag_atom_index(i)
            mag_idx_j = self._get_mag_atom_index(j)
            
            if mag_idx_i is None or mag_idx_j is None:
                continue
            
            J = bond.exchange.matrix
            dl = bond.dl
            
            # Phase factor
            phase = np.exp(2j * np.pi * np.dot(q - self._magstr.k, dl))
            
            # Spin magnitudes
            Si = np.linalg.norm(S[mag_idx_i])
            Sj = np.linalg.norm(S[mag_idx_j])
            
            # Unit vectors
            zi = S[mag_idx_i] / (Si + 1e-10)
            zj = S[mag_idx_j] / (Sj + 1e-10)
            
            # Build local frames
            xi, yi = self._perpendicular_vectors(zi)
            xj, yj = self._perpendicular_vectors(zj)
            
            # Rotate J to local frames
            R_i = np.column_stack([xi, yi, zi])
            R_j = np.column_stack([xj, yj, zj])
            J_local = R_i.T @ J @ R_j
            
            # Prefactor
            pref = np.sqrt(Si * Sj) / 2
            
            # A matrix: (J_xx + J_yy) + i(J_xy - J_yx)
            A[mag_idx_i, mag_idx_j] += pref * phase * (
                (J_local[0, 0] + J_local[1, 1]) + 
                1j * (J_local[0, 1] - J_local[1, 0])
            )
            
            # B matrix: (J_xx - J_yy) + i(J_xy + J_yx)
            B[mag_idx_i, mag_idx_j] += pref * phase * (
                (J_local[0, 0] - J_local[1, 1]) + 
                1j * (J_local[0, 1] + J_local[1, 0])
            )
            
            # Diagonal terms from zz component
            A[mag_idx_i, mag_idx_i] -= Si * J_local[2, 2] * np.dot(zi, zj)
        
        # Single-ion anisotropy contributions
        for atom_idx, A_mat in self._coupling.anisotropy.items():
            mag_idx = self._get_mag_atom_index(atom_idx)
            if mag_idx is None:
                continue
            
            Si = np.linalg.norm(S[mag_idx])
            zi = S[mag_idx] / (Si + 1e-10)
            xi, yi = self._perpendicular_vectors(zi)
            
            R = np.column_stack([xi, yi, zi])
            A_local = R.T @ A_mat @ R
            
            A[mag_idx, mag_idx] += Si * A_local[2, 2]
            A[mag_idx, mag_idx] += (A_local[0, 0] + A_local[1, 1]) / 2
            B[mag_idx, mag_idx] += (A_local[0, 0] - A_local[1, 1]) / 2
            B[mag_idx, mag_idx] += 1j * A_local[0, 1]
        
        # Hermitize A
        A = (A + np.conj(A.T)) / 2
        
        return A, B
    
    def _get_mag_atom_index(self, atom_idx: int) -> Optional[int]:
        """Get the magnetic atom index for a given atom index."""
        mag_idx = 0
        for i, a in enumerate(self._atoms):
            if a.spin > 0:
                if i == atom_idx:
                    return mag_idx
                mag_idx += 1
        return None
    
    def _perpendicular_vectors(self, z: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Get two vectors perpendicular to z."""
        z = z / np.linalg.norm(z)
        
        # Find a vector not parallel to z
        if abs(z[0]) < 0.9:
            ref = np.array([1, 0, 0])
        else:
            ref = np.array([0, 1, 0])
        
        x = np.cross(z, ref)
        x = x / np.linalg.norm(x)
        y = np.cross(z, x)
        
        return x, y
    
    def _calculate_correlation(self, q: np.ndarray, V: np.ndarray,
                                S: np.ndarray, n_mag: int) -> np.ndarray:
        """Calculate spin-spin correlation tensor S^{alpha,beta}(q, omega)."""
        # Simplified calculation - full implementation would use 
        # transformation matrix V to compute correlation functions
        
        Sab = np.zeros((3, 3, n_mag), dtype=np.complex128)
        
        # For now, return approximate isotropic correlation
        for mode in range(n_mag):
            Sab[0, 0, mode] = 1.0
            Sab[1, 1, mode] = 1.0
            Sab[2, 2, mode] = 1.0
        
        return Sab
    
    def _calculate_formfactor(self, Q: np.ndarray) -> np.ndarray:
        """Calculate magnetic form factor for all Q-points."""
        # Simplified implementation
        Q_mag = np.linalg.norm(Q, axis=1)
        # Approximate form factor decay
        ff = np.exp(-0.1 * Q_mag**2)
        return ff
    
    # ========== Utility methods ==========
    
    def _clear_cache(self):
        """Clear cached data."""
        self._cache = {}
    
    def __repr__(self) -> str:
        return (f"SpinW(atoms={self.n_atoms}, mag_atoms={self.n_mag_atoms}, "
                f"bonds={len(self._coupling.bonds)})")


# ========== Convenience functions ==========

def sw_model(name: str, *args, **kwargs) -> SpinW:
    """
    Create a predefined spin wave model.
    
    Parameters
    ----------
    name : str
        Model name: 'chain', 'square', 'triAF', 'squareAF', 'kagome'
    *args : 
        Model-specific parameters (typically J value)
        
    Returns
    -------
    sw : SpinW
        Configured SpinW object
    """
    name = name.lower()
    
    J = args[0] if args else 1.0
    
    if name == 'chain' or name == 'afmchain':
        sw = SpinW()
        sw.genlattice(lat_const=[3, 8, 8], angled=[90, 90, 90], spgr='P 1')
        sw.addatom(r=[0, 0, 0], S=1, label='M1')
        sw.gencoupling(max_distance=4)
        sw.addmatrix(label='J1', value=J)
        sw.addcoupling(mat='J1', bond=1)
        sw.genmagstr(mode='direct', k=[0, 0, 0], 
                     S=np.array([[0, 0, 1]]))
        return sw
    
    elif name == 'square' or name == 'squareaf':
        sw = SpinW()
        sw.genlattice(lat_const=[3, 3, 6], angled=[90, 90, 90], spgr='P 1')
        sw.addatom(r=[0, 0, 0], S=1, label='Cu1')
        sw.gencoupling(max_distance=4)
        sw.addmatrix(label='J1', value=J)
        sw.addcoupling(mat='J1', bond=1)
        sw.genmagstr(mode='direct', k=[0.5, 0.5, 0], 
                     S=np.array([[0, 0, 1]]))
        return sw
    
    elif name == 'triaf' or name == 'triangular':
        sw = SpinW()
        sw.genlattice(lat_const=[3, 3, 6], angled=[90, 90, 120], spgr='P 1')
        sw.addatom(r=[0, 0, 0], S=1, label='Cu1')
        sw.gencoupling(max_distance=4)
        sw.addmatrix(label='J1', value=J)
        sw.addcoupling(mat='J1', bond=1)
        # 120 degree structure
        S = np.array([[1, 0, 0]])
        sw.genmagstr(mode='helical', k=[1/3, 1/3, 0], S=S)
        return sw
    
    elif name == 'kagome':
        sw = SpinW()
        sw.genlattice(lat_const=[6, 6, 6], angled=[90, 90, 120], spgr='P -3')
        sw.addatom(r=[0.5, 0, 0], S=1, label='M1')
        sw.gencoupling(max_distance=4)
        sw.addmatrix(label='J1', value=J)
        sw.addcoupling(mat='J1', bond=1)
        # Coplanar structure
        from .magnetic import generate_120_structure
        magstr = generate_120_structure(3)
        sw._magstr = magstr
        return sw
    
    else:
        raise ValueError(f"Unknown model: {name}")
