"""
Crystal structure handling for DiffTB.

Supports reading from common formats and defines structures
for common materials like Bi2Te3 and MnBi2Te4.
"""

from __future__ import annotations

import re
from pathlib import Path

import jax.numpy as jnp
from beartype.typing import List, NamedTuple, Optional
from jax import Array


class CrystalStructure(NamedTuple):
    """
    Crystal structure specification.

    Attributes
    ----------
        lattice_vectors: Lattice vectors as rows, shape (3, 3), in Angstroms
        positions: Atomic positions in fractional coordinates, shape (n_atoms, 3)
        elements: Element symbol for each atom
        name: Optional structure name
    """

    lattice_vectors: Array
    positions: Array
    elements: List[str]
    name: Optional[str] = None

    @property
    def n_atoms(self) -> int:
        return len(self.elements)

    @property
    def positions_cartesian(self) -> Array:
        """Atomic positions in Cartesian coordinates (Angstroms)."""
        return self.positions @ self.lattice_vectors

    @property
    def volume(self) -> float:
        """Unit cell volume in Å³."""
        return float(jnp.abs(jnp.linalg.det(self.lattice_vectors)))


# =============================================================================
# File I/O
# =============================================================================


def load_structure(filepath: str, format: str = "auto") -> CrystalStructure:
    """
    Load crystal structure from file.

    Args:
        filepath: Path to structure file
        format: File format ('auto', 'vasp', 'xyz', 'cif')

    Returns
    -------
        CrystalStructure object
    """
    path = Path(filepath)

    if format == "auto":
        suffix = path.suffix.lower()
        if suffix in [".vasp", ".poscar"]:
            format = "vasp"
        elif suffix == ".xyz":
            format = "xyz"
        elif suffix == ".cif":
            format = "cif"
        else:
            raise ValueError(
                f"Cannot determine format from extension: {suffix}"
            )

    if format == "vasp":
        return read_poscar(filepath)
    if format == "xyz":
        return read_xyz(filepath)
    raise NotImplementedError(f"Format {format} not yet implemented")


def read_poscar(filepath: str) -> CrystalStructure:
    """Read VASP POSCAR/CONTCAR format."""
    with open(filepath) as f:
        lines = f.readlines()

    name = lines[0].strip()
    scale = float(lines[1])

    # Lattice vectors
    lattice = []
    for i in range(2, 5):
        vec = [float(x) for x in lines[i].split()]
        lattice.append(vec)
    lattice = jnp.array(lattice) * scale

    # Elements and counts
    elements_line = lines[5].split()
    counts_line = [int(x) for x in lines[6].split()]

    elements = []
    for elem, count in zip(elements_line, counts_line, strict=False):
        elements.extend([elem] * count)

    # Check for selective dynamics / direct/cartesian
    line_idx = 7
    if lines[line_idx][0].lower() in ["s", "selective"]:
        line_idx += 1

    coord_type = lines[line_idx][0].lower()
    line_idx += 1

    # Positions
    n_atoms = sum(counts_line)
    positions = []
    for i in range(n_atoms):
        pos = [float(x) for x in lines[line_idx + i].split()[:3]]
        positions.append(pos)
    positions = jnp.array(positions)

    # Convert to fractional if Cartesian
    if coord_type == "c":
        positions = positions @ jnp.linalg.inv(lattice)

    return CrystalStructure(
        lattice_vectors=lattice,
        positions=positions,
        elements=elements,
        name=name,
    )


def read_xyz(filepath: str) -> CrystalStructure:
    """
    Read XYZ format (assumes lattice in comment line).

    Extended XYZ format: lattice vectors in comment as
    Lattice="ax ay az bx by bz cx cy cz"
    """
    with open(filepath) as f:
        lines = f.readlines()

    n_atoms = int(lines[0].strip())
    comment = lines[1].strip()

    # Parse lattice from comment
    lattice = jnp.eye(3) * 20.0  # Default: large box
    if "Lattice=" in comment:
        match = re.search(r'Lattice="([^"]+)"', comment)
        if match:
            vals = [float(x) for x in match.group(1).split()]
            lattice = jnp.array(vals).reshape(3, 3)

    # Parse atoms
    elements = []
    positions = []
    for i in range(n_atoms):
        parts = lines[2 + i].split()
        elements.append(parts[0])
        pos = [float(x) for x in parts[1:4]]
        positions.append(pos)

    positions = jnp.array(positions)

    # Convert Cartesian to fractional
    positions_frac = positions @ jnp.linalg.inv(lattice)

    return CrystalStructure(
        lattice_vectors=lattice, positions=positions_frac, elements=elements
    )


def write_poscar(structure: CrystalStructure, filepath: str) -> None:
    """Write structure to VASP POSCAR format."""
    with open(filepath, "w") as f:
        f.write(f"{structure.name or 'DiffTB structure'}\n")
        f.write("1.0\n")

        for vec in structure.lattice_vectors:
            f.write(f"  {vec[0]:16.10f} {vec[1]:16.10f} {vec[2]:16.10f}\n")

        # Count elements
        unique_elements = []
        counts = []
        for elem in structure.elements:
            if not unique_elements or elem != unique_elements[-1]:
                unique_elements.append(elem)
                counts.append(1)
            else:
                counts[-1] += 1

        f.write("  " + "  ".join(unique_elements) + "\n")
        f.write("  " + "  ".join(map(str, counts)) + "\n")
        f.write("Direct\n")

        for pos in structure.positions:
            f.write(f"  {pos[0]:16.10f} {pos[1]:16.10f} {pos[2]:16.10f}\n")


# =============================================================================
# Built-in structures for testing
# =============================================================================


def bi2te3_structure() -> CrystalStructure:
    """
    Create Bi2Te3 crystal structure (rhombohedral, R-3m).

    Quintuple layer structure: Te1-Bi-Te2-Bi-Te1

    Lattice parameters from experiment:
    a = 4.384 Å, c = 30.487 Å (hexagonal setting)
    """
    # Hexagonal setting
    a = 4.384
    c = 30.487

    # Lattice vectors
    lattice = jnp.array(
        [[a, 0, 0], [-a / 2, a * jnp.sqrt(3) / 2, 0], [0, 0, c]]
    )

    # Atomic positions (fractional, hexagonal setting)
    # 1 formula unit in primitive cell, 3 in hexagonal
    positions = jnp.array(
        [
            # Quintuple layer 1
            [0.0, 0.0, 0.0],  # Te1
            [0.0, 0.0, 0.2097],  # Bi
            [0.0, 0.0, 0.4000],  # Te2
            [0.0, 0.0, 0.5903],  # Bi
            [0.0, 0.0, 0.8000],  # Te1
        ]
    )

    elements = ["Te", "Bi", "Te", "Bi", "Te"]

    return CrystalStructure(
        lattice_vectors=lattice,
        positions=positions,
        elements=elements,
        name="Bi2Te3",
    )


def mnbi2te4_structure() -> CrystalStructure:
    """
    Create MnBi2Te4 crystal structure (rhombohedral, R-3m).

    Septuple layer structure: Te-Bi-Te-Mn-Te-Bi-Te

    Lattice parameters from experiment:
    a = 4.334 Å, c = 40.91 Å (hexagonal setting)
    """
    a = 4.334
    c = 40.91

    lattice = jnp.array(
        [[a, 0, 0], [-a / 2, a * jnp.sqrt(3) / 2, 0], [0, 0, c]]
    )

    # Atomic positions for septuple layer
    positions = jnp.array(
        [
            [0.0, 0.0, 0.0],  # Te (outer)
            [0.0, 0.0, 0.0863],  # Bi
            [0.0, 0.0, 0.1715],  # Te (inner)
            [0.0, 0.0, 0.25],  # Mn (center)
            [0.0, 0.0, 0.3285],  # Te (inner)
            [0.0, 0.0, 0.4137],  # Bi
            [0.0, 0.0, 0.5],  # Te (outer)
        ]
    )

    elements = ["Te", "Bi", "Te", "Mn", "Te", "Bi", "Te"]

    return CrystalStructure(
        lattice_vectors=lattice,
        positions=positions,
        elements=elements,
        name="MnBi2Te4",
    )


def graphene_structure(a: float = 2.46) -> CrystalStructure:
    """
    Create graphene structure.

    Args:
        a: Lattice constant in Angstroms (default 2.46)
    """
    lattice = jnp.array(
        [[a, 0, 0], [-a / 2, a * jnp.sqrt(3) / 2, 0], [0, 0, 20.0]]  # Vacuum
    )

    positions = jnp.array([[0.0, 0.0, 0.5], [1 / 3, 2 / 3, 0.5]])

    elements = ["C", "C"]

    return CrystalStructure(
        lattice_vectors=lattice,
        positions=positions,
        elements=elements,
        name="Graphene",
    )


# =============================================================================
# Structure manipulation
# =============================================================================


def make_supercell(
    structure: CrystalStructure, n1: int, n2: int, n3: int
) -> CrystalStructure:
    """
    Create a supercell by repeating the unit cell.

    Args:
        structure: Original structure
        n1, n2, n3: Repetitions along each lattice vector

    Returns
    -------
        Supercell structure
    """
    new_lattice = structure.lattice_vectors * jnp.array([[n1], [n2], [n3]])

    new_positions = []
    new_elements = []

    for i1 in range(n1):
        for i2 in range(n2):
            for i3 in range(n3):
                shift = jnp.array([i1 / n1, i2 / n2, i3 / n3])
                for pos, elem in zip(
                    structure.positions, structure.elements, strict=False
                ):
                    new_pos = (pos + shift) / jnp.array([n1, n2, n3])
                    new_positions.append(new_pos)
                    new_elements.append(elem)

    return CrystalStructure(
        lattice_vectors=new_lattice,
        positions=jnp.stack(new_positions),
        elements=new_elements,
        name=f"{structure.name} {n1}x{n2}x{n3}",
    )


def make_slab(
    structure: CrystalStructure,
    n_layers: int,
    vacuum: float = 15.0,
    surface_normal: int = 2,
) -> CrystalStructure:
    """
    Create a slab from bulk structure.

    Args:
        structure: Bulk structure
        n_layers: Number of layers
        vacuum: Vacuum thickness in Angstroms
        surface_normal: Direction of surface normal (0, 1, or 2)

    Returns
    -------
        Slab structure
    """
    # This is simplified - proper slab construction
    # requires identifying layer structure

    # Repeat in surface normal direction
    repeats = [1, 1, 1]
    repeats[surface_normal] = n_layers
    slab = make_supercell(structure, *repeats)

    # Add vacuum
    new_lattice = slab.lattice_vectors.at[surface_normal, surface_normal].add(
        vacuum
    )

    # Shift atoms to center of cell
    shift = vacuum / (2 * new_lattice[surface_normal, surface_normal])
    new_positions = slab.positions.at[:, surface_normal].add(shift)

    return CrystalStructure(
        lattice_vectors=new_lattice,
        positions=new_positions,
        elements=slab.elements,
        name=f"{structure.name} slab ({n_layers} layers)",
    )
