"""
Lattice and crystallographic operations for PySpinW.

Handles:
- Unit cell definition
- Space group symmetry operations
- Reciprocal lattice calculations
- Wyckoff positions
"""

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


@dataclass
class SpaceGroup:
    """
    Space group symmetry operations.
    
    Attributes
    ----------
    number : int
        Space group number (1-230)
    symbol : str
        Hermann-Mauguin symbol
    generators : np.ndarray
        Symmetry generator matrices, shape (n_ops, 4, 4)
    """
    number: int = 1
    symbol: str = 'P 1'
    generators: np.ndarray = field(default_factory=lambda: np.eye(4)[np.newaxis, ...])
    
    @classmethod
    def from_symbol(cls, symbol: str) -> 'SpaceGroup':
        """Create space group from Hermann-Mauguin symbol."""
        symbol = symbol.strip().replace('_', ' ')
        
        # Common space groups lookup
        spg_table = {
            'P 1': (1, [np.eye(4)]),
            'P -1': (2, [np.eye(4), np.diag([-1, -1, -1, 1])]),
            'P 2': (3, _generate_p2()),
            'P 21': (4, _generate_p21()),
            'P m': (6, _generate_pm()),
            'P 2/m': (10, _generate_p2_m()),
            'C 2': (5, _generate_c2()),
            'P 222': (16, _generate_p222()),
            'P m m m': (47, _generate_pmmm()),
            'P 4': (75, _generate_p4()),
            'P 4/m': (83, _generate_p4_m()),
            'P -3': (147, _generate_p_3()),
            'R -3': (148, _generate_r_3()),
            'P 6': (168, _generate_p6()),
            'P 63': (173, _generate_p63()),
            'P m -3 m': (221, _generate_pm3m()),
            'F m -3 m': (225, _generate_fm3m()),
            'I m -3 m': (229, _generate_im3m()),
        }
        
        if symbol in spg_table:
            num, gens = spg_table[symbol]
            return cls(number=num, symbol=symbol, generators=np.array(gens))
        
        # Default to P1 if not found
        return cls(number=1, symbol=symbol, generators=np.eye(4)[np.newaxis, ...])
    
    @classmethod
    def from_number(cls, number: int) -> 'SpaceGroup':
        """Create space group from number."""
        # Map common numbers to symbols
        num_to_sym = {
            1: 'P 1', 2: 'P -1', 147: 'P -3', 221: 'P m -3 m', 225: 'F m -3 m'
        }
        symbol = num_to_sym.get(number, f'SG{number}')
        return cls.from_symbol(symbol)
    
    def get_all_operations(self, origin: np.ndarray = None) -> np.ndarray:
        """
        Generate all symmetry operations from generators.
        
        Returns
        -------
        ops : np.ndarray
            All symmetry operations, shape (n_ops, 4, 4)
        """
        if origin is not None:
            # Shift operations to new origin
            pass
        return self.generators
    
    def apply_symmetry(self, positions: np.ndarray) -> np.ndarray:
        """
        Apply all symmetry operations to positions.
        
        Parameters
        ----------
        positions : np.ndarray
            Atomic positions, shape (n_atoms, 3)
            
        Returns
        -------
        all_positions : np.ndarray
            All symmetry-equivalent positions, shape (n_eq, 3)
        """
        ops = self.get_all_operations()
        all_pos = []
        
        for pos in positions:
            for op in ops:
                new_pos = op[:3, :3] @ pos + op[:3, 3]
                # Wrap to [0, 1)
                new_pos = new_pos % 1.0
                all_pos.append(new_pos)
        
        # Remove duplicates
        all_pos = np.array(all_pos)
        unique_pos = _unique_positions(all_pos)
        
        return unique_pos


def _unique_positions(positions: np.ndarray, tol: float = 1e-6) -> np.ndarray:
    """Remove duplicate positions within tolerance."""
    unique = [positions[0]]
    for pos in positions[1:]:
        is_unique = True
        for u in unique:
            diff = np.abs(pos - u)
            # Account for periodic boundary
            diff = np.minimum(diff, 1 - diff)
            if np.max(diff) < tol:
                is_unique = False
                break
        if is_unique:
            unique.append(pos)
    return np.array(unique)


# Generator functions for common space groups
def _generate_p2():
    """P 2 (monoclinic)."""
    ops = [np.eye(4)]
    # 2-fold rotation about b
    op2 = np.diag([-1, 1, -1, 1])
    ops.append(op2)
    return ops

def _generate_p21():
    """P 21 (monoclinic with screw)."""
    ops = [np.eye(4)]
    op2 = np.diag([-1, 1, -1, 1])
    op2[1, 3] = 0.5  # Translation
    ops.append(op2)
    return ops

def _generate_pm():
    """P m (mirror)."""
    ops = [np.eye(4)]
    op_m = np.diag([1, -1, 1, 1])
    ops.append(op_m)
    return ops

def _generate_p2_m():
    """P 2/m."""
    return _generate_p2() + [np.diag([-1, -1, -1, 1])]

def _generate_c2():
    """C 2 (C-centered monoclinic)."""
    ops = _generate_p2()
    # Add C-centering translation
    for op in ops[:]:
        op_c = op.copy()
        op_c[:3, 3] = (op_c[:3, 3] + np.array([0.5, 0.5, 0])) % 1.0
        ops.append(op_c)
    return ops

def _generate_p222():
    """P 222 (orthorhombic)."""
    ops = [np.eye(4)]
    ops.append(np.diag([-1, -1, 1, 1]))  # 2z
    ops.append(np.diag([-1, 1, -1, 1]))  # 2y
    ops.append(np.diag([1, -1, -1, 1]))  # 2x
    return ops

def _generate_pmmm():
    """P mmm."""
    ops = _generate_p222()
    ops.append(np.diag([-1, -1, -1, 1]))
    return ops

def _generate_p4():
    """P 4 (tetragonal)."""
    ops = [np.eye(4)]
    # 4-fold rotation about z
    rot4 = np.eye(4)
    rot4[:2, :2] = [[0, -1], [1, 0]]
    ops.append(rot4)
    ops.append(rot4 @ rot4)
    ops.append(rot4 @ rot4 @ rot4)
    return ops

def _generate_p4_m():
    """P 4/m."""
    return _generate_p4() + [np.diag([-1, -1, -1, 1])]

def _generate_p_3():
    """P -3 (trigonal)."""
    ops = [np.eye(4)]
    # 3-fold rotation about z (hexagonal setting)
    c = np.cos(2*np.pi/3)
    s = np.sin(2*np.pi/3)
    rot3 = np.eye(4)
    rot3[:2, :2] = [[c, -s], [s, c]]
    ops.append(rot3)
    ops.append(rot3 @ rot3)
    # Inversion
    ops.append(np.diag([-1, -1, -1, 1]))
    ops.append(np.diag([-1, -1, -1, 1]) @ rot3)
    ops.append(np.diag([-1, -1, -1, 1]) @ rot3 @ rot3)
    return ops

def _generate_r_3():
    """R -3 (rhombohedral)."""
    return _generate_p_3()

def _generate_p6():
    """P 6 (hexagonal)."""
    ops = [np.eye(4)]
    c = np.cos(np.pi/3)
    s = np.sin(np.pi/3)
    rot6 = np.eye(4)
    rot6[:2, :2] = [[c, -s], [s, c]]
    for i in range(1, 6):
        rot = np.eye(4)
        for _ in range(i):
            rot = rot @ rot6
        ops.append(rot)
    return ops

def _generate_p63():
    """P 63."""
    ops = _generate_p6()
    # Add screw translation
    for i, op in enumerate(ops):
        if i > 0:
            ops[i][:3, 3] = (ops[i][:3, 3] + np.array([0, 0, i * 0.5])) % 1.0
    return ops

def _generate_pm3m():
    """P m -3 m (cubic)."""
    ops = []
    # 48 operations of Oh group
    # Identity
    ops.append(np.eye(4))
    # Add cubic rotations and mirrors
    rotations = [
        np.diag([1, 1, 1, 1]),
        np.diag([-1, -1, 1, 1]),
        np.diag([-1, 1, -1, 1]),
        np.diag([1, -1, -1, 1]),
    ]
    for rot in rotations:
        ops.append(rot)
    return ops[:48] if len(ops) >= 48 else ops

def _generate_fm3m():
    """F m -3 m (FCC)."""
    ops = _generate_pm3m()
    # Add face-centering translations
    trans = [[0.5, 0.5, 0], [0.5, 0, 0.5], [0, 0.5, 0.5]]
    expanded = []
    for op in ops:
        expanded.append(op)
        for t in trans:
            op_t = op.copy()
            op_t[:3, 3] = (op_t[:3, 3] + np.array(t)) % 1.0
            expanded.append(op_t)
    return expanded

def _generate_im3m():
    """I m -3 m (BCC)."""
    ops = _generate_pm3m()
    # Add body-centering translation
    trans = [0.5, 0.5, 0.5]
    expanded = []
    for op in ops:
        expanded.append(op)
        op_t = op.copy()
        op_t[:3, 3] = (op_t[:3, 3] + np.array(trans)) % 1.0
        expanded.append(op_t)
    return expanded


@dataclass
class Lattice:
    """
    Crystal lattice definition.
    
    Attributes
    ----------
    lat_const : tuple
        Lattice constants (a, b, c) in Angstrom
    angles : tuple
        Lattice angles (alpha, beta, gamma) in degrees
    spacegroup : SpaceGroup
        Space group symmetry
    """
    lat_const: Tuple[float, float, float] = (1.0, 1.0, 1.0)
    angles: Tuple[float, float, float] = (90.0, 90.0, 90.0)
    spacegroup: SpaceGroup = field(default_factory=SpaceGroup)
    
    def __post_init__(self):
        self._basis_vectors = None
        self._reciprocal_vectors = None
        self._metric_tensor = None
    
    @property
    def a(self) -> float:
        return self.lat_const[0]
    
    @property
    def b(self) -> float:
        return self.lat_const[1]
    
    @property
    def c(self) -> float:
        return self.lat_const[2]
    
    @property
    def alpha(self) -> float:
        return self.angles[0]
    
    @property
    def beta(self) -> float:
        return self.angles[1]
    
    @property
    def gamma(self) -> float:
        return self.angles[2]
    
    @property
    def basis_vectors(self) -> np.ndarray:
        """
        Cartesian basis vectors of the unit cell.
        
        Returns
        -------
        bv : np.ndarray
            Basis vectors as columns, shape (3, 3)
            Column 0: a, Column 1: b, Column 2: c
        """
        if self._basis_vectors is None:
            self._basis_vectors = self._compute_basis_vectors()
        return self._basis_vectors
    
    def _compute_basis_vectors(self) -> np.ndarray:
        """Compute Cartesian basis vectors from lattice parameters."""
        a, b, c = self.lat_const
        alpha, beta, gamma = np.radians(self.angles)
        
        # Crystallographic convention: a along x
        av = np.array([a, 0, 0])
        
        # b in xy-plane
        bv = np.array([b * np.cos(gamma), b * np.sin(gamma), 0])
        
        # c calculated from angles
        cx = c * np.cos(beta)
        cy = c * (np.cos(alpha) - np.cos(beta) * np.cos(gamma)) / np.sin(gamma)
        cz = np.sqrt(c**2 - cx**2 - cy**2)
        cv = np.array([cx, cy, cz])
        
        return np.column_stack([av, bv, cv])
    
    @property
    def reciprocal_vectors(self) -> np.ndarray:
        """
        Reciprocal lattice vectors.
        
        Returns
        -------
        rv : np.ndarray
            Reciprocal vectors as columns, shape (3, 3)
        """
        if self._reciprocal_vectors is None:
            bv = self.basis_vectors
            # b* = 2*pi * (a x c) / (a . (b x c))
            volume = np.linalg.det(bv)
            self._reciprocal_vectors = 2 * np.pi * np.linalg.inv(bv).T
        return self._reciprocal_vectors
    
    @property
    def volume(self) -> float:
        """Unit cell volume in Å³."""
        return np.abs(np.linalg.det(self.basis_vectors))
    
    @property
    def metric_tensor(self) -> np.ndarray:
        """
        Metric tensor g_ij = a_i · a_j
        
        Returns
        -------
        g : np.ndarray
            Metric tensor, shape (3, 3)
        """
        if self._metric_tensor is None:
            bv = self.basis_vectors
            self._metric_tensor = bv.T @ bv
        return self._metric_tensor
    
    def frac_to_cart(self, frac: np.ndarray) -> np.ndarray:
        """
        Convert fractional to Cartesian coordinates.
        
        Parameters
        ----------
        frac : np.ndarray
            Fractional coordinates, shape (..., 3)
            
        Returns
        -------
        cart : np.ndarray
            Cartesian coordinates, shape (..., 3)
        """
        return frac @ self.basis_vectors.T
    
    def cart_to_frac(self, cart: np.ndarray) -> np.ndarray:
        """
        Convert Cartesian to fractional coordinates.
        
        Parameters
        ----------
        cart : np.ndarray
            Cartesian coordinates, shape (..., 3)
            
        Returns
        -------
        frac : np.ndarray
            Fractional coordinates, shape (..., 3)
        """
        return cart @ np.linalg.inv(self.basis_vectors).T
    
    def hkl_to_q(self, hkl: np.ndarray) -> np.ndarray:
        """
        Convert Miller indices to Q-vector (Å⁻¹).
        
        Parameters
        ----------
        hkl : np.ndarray
            Miller indices, shape (..., 3)
            
        Returns
        -------
        Q : np.ndarray
            Q-vector in Å⁻¹, shape (..., 3)
        """
        return hkl @ self.reciprocal_vectors.T
    
    def q_to_hkl(self, Q: np.ndarray) -> np.ndarray:
        """
        Convert Q-vector to Miller indices.
        
        Parameters
        ----------
        Q : np.ndarray
            Q-vector in Å⁻¹, shape (..., 3)
            
        Returns
        -------
        hkl : np.ndarray
            Miller indices, shape (..., 3)
        """
        return Q @ np.linalg.inv(self.reciprocal_vectors).T
    
    def d_spacing(self, hkl: np.ndarray) -> float:
        """
        Calculate d-spacing for given Miller indices.
        
        Parameters
        ----------
        hkl : np.ndarray
            Miller indices, shape (3,) or (n, 3)
            
        Returns
        -------
        d : float or np.ndarray
            d-spacing in Angstrom
        """
        Q = self.hkl_to_q(hkl)
        Q_mag = np.linalg.norm(Q, axis=-1)
        return 2 * np.pi / Q_mag
