"""
Magnetic coupling and exchange interactions for PySpinW.

Handles:
- Exchange coupling matrices (Heisenberg, anisotropic, DM)
- Bond generation
- Single-ion anisotropy
- External magnetic field
"""

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


@dataclass
class ExchangeMatrix:
    """
    3x3 exchange coupling matrix.
    
    Can represent:
    - Isotropic Heisenberg: J * I
    - Anisotropic exchange: diag(Jx, Jy, Jz)
    - Dzyaloshinskii-Moriya: antisymmetric part
    - General tensor exchange
    
    Attributes
    ----------
    matrix : np.ndarray
        3x3 exchange matrix in meV
    label : str
        Label for this interaction (e.g., 'J1', 'J2')
    color : str
        Color for visualization
    """
    matrix: np.ndarray = field(default_factory=lambda: np.eye(3))
    label: str = 'J'
    color: str = 'gray'
    
    def __post_init__(self):
        self.matrix = np.asarray(self.matrix, dtype=np.float64)
        if self.matrix.shape == ():
            # Scalar -> isotropic
            self.matrix = float(self.matrix) * np.eye(3)
        elif self.matrix.shape == (3,):
            # Diagonal
            self.matrix = np.diag(self.matrix)
    
    @classmethod
    def heisenberg(cls, J: float, label: str = 'J') -> 'ExchangeMatrix':
        """Create isotropic Heisenberg exchange."""
        return cls(matrix=J * np.eye(3), label=label)
    
    @classmethod
    def anisotropic(cls, Jx: float, Jy: float, Jz: float, 
                    label: str = 'J') -> 'ExchangeMatrix':
        """Create diagonal anisotropic exchange."""
        return cls(matrix=np.diag([Jx, Jy, Jz]), label=label)
    
    @classmethod
    def dzyaloshinskii_moriya(cls, D: np.ndarray, 
                              label: str = 'DM') -> 'ExchangeMatrix':
        """
        Create Dzyaloshinskii-Moriya interaction.
        
        The DM interaction is: D · (S_i × S_j)
        
        Parameters
        ----------
        D : np.ndarray
            DM vector [Dx, Dy, Dz]
        """
        D = np.asarray(D)
        matrix = np.array([
            [0, -D[2], D[1]],
            [D[2], 0, -D[0]],
            [-D[1], D[0], 0]
        ])
        return cls(matrix=matrix, label=label)
    
    @property
    def isotropic_part(self) -> float:
        """Isotropic (trace) component."""
        return np.trace(self.matrix) / 3
    
    @property
    def symmetric_part(self) -> np.ndarray:
        """Symmetric part of exchange matrix."""
        return (self.matrix + self.matrix.T) / 2
    
    @property
    def antisymmetric_part(self) -> np.ndarray:
        """Antisymmetric part (DM interaction)."""
        return (self.matrix - self.matrix.T) / 2
    
    @property
    def dm_vector(self) -> np.ndarray:
        """Extract DM vector from antisymmetric part."""
        A = self.antisymmetric_part
        # For DM matrix: A_ij = epsilon_ijk * D_k
        # So: D = [A[2,1], A[0,2], A[1,0]]
        return np.array([A[2, 1], A[0, 2], A[1, 0]])
    
    @property 
    def anisotropy_axes(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get principal axes and values of symmetric exchange."""
        sym = self.symmetric_part
        eigenvalues, eigenvectors = np.linalg.eigh(sym)
        return eigenvalues, eigenvectors


@dataclass
class Bond:
    """
    Magnetic bond between two atoms.
    
    Attributes
    ----------
    atom1 : int
        Index of first atom
    atom2 : int
        Index of second atom
    dl : np.ndarray
        Lattice vector connecting cells [da, db, dc]
    distance : float
        Bond distance in Angstrom
    exchange : ExchangeMatrix
        Exchange coupling for this bond
    idx : int
        Bond type index (bonds of same type have same exchange)
    """
    atom1: int = 0
    atom2: int = 0
    dl: np.ndarray = field(default_factory=lambda: np.zeros(3, dtype=int))
    distance: float = 0.0
    exchange: ExchangeMatrix = None
    idx: int = 0
    
    def __post_init__(self):
        self.dl = np.asarray(self.dl, dtype=int)


@dataclass
class Coupling:
    """
    Collection of magnetic couplings and interactions.
    
    Manages all magnetic interactions in the system:
    - Pair exchange interactions
    - Single-ion anisotropy
    - External magnetic field
    - Biquadratic exchange
    """
    bonds: List[Bond] = field(default_factory=list)
    exchange_matrices: Dict[str, ExchangeMatrix] = field(default_factory=dict)
    
    # Single-ion anisotropy per atom type
    anisotropy: Dict[int, np.ndarray] = field(default_factory=dict)
    
    # External magnetic field [Bx, By, Bz] in Tesla
    magnetic_field: np.ndarray = field(default_factory=lambda: np.zeros(3))
    
    # Temperature in Kelvin
    temperature: float = 0.0
    
    def add_matrix(self, label: str, value: Union[float, np.ndarray], 
                   color: str = 'gray') -> ExchangeMatrix:
        """
        Add an exchange matrix.
        
        Parameters
        ----------
        label : str
            Matrix label (e.g., 'J1', 'J2', 'DM')
        value : float or np.ndarray
            Scalar for isotropic, 3-vector for DM, 3x3 for general
        color : str
            Visualization color
        """
        value = np.asarray(value)
        
        if value.shape == ():
            mat = ExchangeMatrix.heisenberg(float(value), label)
        elif value.shape == (3,):
            # Check if it's meant to be DM or diagonal anisotropic
            if 'DM' in label.upper() or 'D' == label[0].upper():
                mat = ExchangeMatrix.dzyaloshinskii_moriya(value, label)
            else:
                mat = ExchangeMatrix.anisotropic(*value, label)
        else:
            mat = ExchangeMatrix(matrix=value, label=label)
        
        mat.color = color
        self.exchange_matrices[label] = mat
        return mat
    
    def add_coupling(self, mat: str, bond_idx: int):
        """
        Assign an exchange matrix to bonds of given type.
        
        Parameters
        ----------
        mat : str
            Label of exchange matrix
        bond_idx : int
            Bond type index (from gencoupling)
        """
        if mat not in self.exchange_matrices:
            raise ValueError(f"Exchange matrix '{mat}' not defined")
        
        exchange = self.exchange_matrices[mat]
        for bond in self.bonds:
            if bond.idx == bond_idx:
                bond.exchange = exchange
    
    def add_anisotropy(self, mat: str, atom_idx: int):
        """
        Add single-ion anisotropy to an atom.
        
        Parameters
        ----------
        mat : str
            Label of anisotropy matrix
        atom_idx : int
            Atom index
        """
        if mat not in self.exchange_matrices:
            raise ValueError(f"Matrix '{mat}' not defined")
        
        self.anisotropy[atom_idx] = self.exchange_matrices[mat].matrix
    
    def set_field(self, B: np.ndarray):
        """Set external magnetic field in Tesla."""
        self.magnetic_field = np.asarray(B)
    
    def set_temperature(self, T: float):
        """Set temperature in Kelvin."""
        self.temperature = T
    
    def generate_bonds(self, positions: np.ndarray, lattice, 
                       max_distance: float) -> List[Bond]:
        """
        Generate all bonds within a cutoff distance.
        
        Parameters
        ----------
        positions : np.ndarray
            Atomic positions in fractional coordinates, shape (n, 3)
        lattice : Lattice
            Crystal lattice
        max_distance : float
            Maximum bond distance in Angstrom
            
        Returns
        -------
        bonds : List[Bond]
            Generated bonds sorted by distance
        """
        n_atoms = len(positions)
        self.bonds = []
        
        # Search over neighboring cells
        search_range = int(np.ceil(max_distance / np.min(lattice.lat_const))) + 1
        
        bond_types = {}  # (distance, atom_pair_type) -> bond_idx
        
        for i in range(n_atoms):
            for j in range(n_atoms):
                for da in range(-search_range, search_range + 1):
                    for db in range(-search_range, search_range + 1):
                        for dc in range(-search_range, search_range + 1):
                            if i == j and da == 0 and db == 0 and dc == 0:
                                continue
                            
                            # Vector from atom i to atom j
                            dl = np.array([da, db, dc])
                            dr_frac = positions[j] - positions[i] + dl
                            dr_cart = lattice.frac_to_cart(dr_frac)
                            dist = np.linalg.norm(dr_cart)
                            
                            if dist <= max_distance:
                                # Round distance to handle floating point
                                dist_key = round(dist, 4)
                                pair_key = tuple(sorted([i, j]))
                                type_key = (dist_key, pair_key)
                                
                                if type_key not in bond_types:
                                    bond_types[type_key] = len(bond_types) + 1
                                
                                bond = Bond(
                                    atom1=i,
                                    atom2=j,
                                    dl=dl,
                                    distance=dist,
                                    idx=bond_types[type_key]
                                )
                                self.bonds.append(bond)
        
        # Sort by distance
        self.bonds.sort(key=lambda b: (b.distance, b.idx))
        
        return self.bonds
    
    def get_bonds_by_type(self, bond_idx: int) -> List[Bond]:
        """Get all bonds of a given type."""
        return [b for b in self.bonds if b.idx == bond_idx]
    
    @property
    def n_bond_types(self) -> int:
        """Number of unique bond types."""
        if not self.bonds:
            return 0
        return max(b.idx for b in self.bonds)
    
    def list_bonds(self) -> str:
        """Generate a summary table of bond types."""
        if not self.bonds:
            return "No bonds defined"
        
        lines = ["Bond  Atoms     Distance  Exchange"]
        lines.append("-" * 40)
        
        seen_types = set()
        for bond in self.bonds:
            if bond.idx in seen_types:
                continue
            seen_types.add(bond.idx)
            
            ex_str = bond.exchange.label if bond.exchange else "-"
            lines.append(
                f"{bond.idx:4d}  {bond.atom1}-{bond.atom2}      "
                f"{bond.distance:7.3f}  {ex_str}"
            )
        
        return "\n".join(lines)


def quick_hamiltonian(spins: np.ndarray, coupling: Coupling, 
                      lattice, k: np.ndarray = None, n: np.ndarray = None) -> float:
    """
    Calculate classical Hamiltonian energy.
    
    For structures with propagation vector k ≠ 0, the spin at position R + dl
    is rotated by angle 2π k·dl around the spiral axis n.
    
    For collinear structures (where spins are parallel to n), we use the
    fact that the angle between neighboring spins is simply 2π k·dl.
    
    Parameters
    ----------
    spins : np.ndarray
        Spin vectors, shape (n_atoms, 3)
    coupling : Coupling
        Coupling definitions
    lattice : Lattice
        Crystal lattice
    k : np.ndarray, optional
        Propagation vector (default: [0, 0, 0])
    n : np.ndarray, optional
        Spiral axis (default: [0, 0, 1])
        
    Returns
    -------
    E : float
        Total energy in meV
    """
    E = 0.0
    n_atoms = len(spins)
    
    if k is None:
        k = np.zeros(3)
    k = np.asarray(k)
    
    if n is None:
        n = np.array([0, 0, 1])
    n = np.asarray(n)
    n = n / (np.linalg.norm(n) + 1e-10)
    
    def rotation_matrix(axis, angle):
        """Rodrigues' rotation formula."""
        axis = axis / (np.linalg.norm(axis) + 1e-10)
        K = np.array([
            [0, -axis[2], axis[1]],
            [axis[2], 0, -axis[0]],
            [-axis[1], axis[0], 0]
        ])
        return np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K
    
    # Deduplicate bonds: (i,j,dl) and (j,i,-dl) are the same bond
    seen = set()
    unique_bonds = []
    for bond in coupling.bonds:
        if bond.exchange is None:
            continue
        i, j = bond.atom1, bond.atom2
        dl = tuple(int(x) for x in bond.dl)
        neg_dl = tuple(-x for x in dl)
        
        key1 = (i, j, dl)
        key2 = (j, i, neg_dl)
        
        if key1 not in seen and key2 not in seen:
            seen.add(key1)
            seen.add(key2)
            unique_bonds.append(bond)
    
    # Exchange energy with propagation vector
    for bond in unique_bonds:
        
        S_i = spins[bond.atom1]
        S_j_base = spins[bond.atom2]
        J = bond.exchange.matrix
        dl = np.asarray(bond.dl)
        
        # Rotation angle for propagation vector
        angle = 2 * np.pi * np.dot(k, dl)
        
        # Check if structure is collinear (spins parallel to n)
        S_i_norm = S_i / (np.linalg.norm(S_i) + 1e-10)
        is_collinear = abs(abs(np.dot(S_i_norm, n)) - 1.0) < 0.01
        
        if is_collinear and abs(angle) > 1e-6:
            # For collinear structures with k ≠ 0:
            # Neighboring spin flips sign when angle = π (k·dl = 0.5)
            # Energy = J * |S_i| * |S_j| * cos(angle) for isotropic exchange
            # This gives the AFM energy correctly
            cos_angle = np.cos(angle)
            # For general J matrix, scale the interaction by cos(angle)
            # This is exact for isotropic Heisenberg
            S_j_effective = S_j_base * cos_angle
            E += S_i @ J @ S_j_effective
        else:
            # For non-collinear (helical) structures:
            # Rotate S_j around the spiral axis n
            R = rotation_matrix(n, angle)
            S_j_rotated = R @ S_j_base
            E += S_i @ J @ S_j_rotated
    
    # Single-ion anisotropy
    for atom_idx, A in coupling.anisotropy.items():
        if atom_idx < n_atoms:
            S = spins[atom_idx]
            E += S @ A @ S
    
    # Zeeman energy
    if np.any(coupling.magnetic_field != 0):
        from .constants import CONSTANTS
        for i in range(n_atoms):
            E -= CONSTANTS.muB * spins[i] @ coupling.magnetic_field
    
    return E
