import numpy as np
from dataclasses import dataclass
from typing import List

ATOMIC_NUMBERS = {
    'H': 1, 'He': 2, 'Li': 3, 'Be': 4, 'B': 5,
    'C': 6, 'N': 7, 'O': 8, 'F': 9, 'Ne': 10,
    # Soft pseudopotential variants (same atomic number)
    'H_soft': 1,
}
ANGSTROM_TO_BOHR = 1.8897259886


@dataclass
class Atom:
    symbol: str
    coords: np.ndarray  # Bohr

    @property
    def Z(self) -> int:
        return ATOMIC_NUMBERS[self.symbol]


class Molecule:
    def __init__(self, atoms: List[Atom], charge: int = 0, spin: int = 0):
        self.atoms = atoms
        self.charge = charge
        self.spin = spin

    @property
    def n_electrons(self) -> int:
        return sum(a.Z for a in self.atoms) - self.charge

    @property
    def n_alpha(self) -> int:
        return (self.n_electrons + self.spin) // 2

    @property
    def n_beta(self) -> int:
        return (self.n_electrons - self.spin) // 2

    @property
    def nuclear_repulsion(self) -> float:
        E = 0.0
        for i, a in enumerate(self.atoms):
            for j in range(i + 1, len(self.atoms)):
                b = self.atoms[j]
                R = np.linalg.norm(a.coords - b.coords)
                E += a.Z * b.Z / R
        return E

    @classmethod
    def from_xyz(cls, xyz: str, charge: int = 0, spin: int = 0, unit: str = 'angstrom'):
        atoms = []
        for line in xyz.strip().splitlines():
            parts = line.split()
            if len(parts) < 4:
                continue
            coords = np.array([float(x) for x in parts[1:4]])
            if unit == 'angstrom':
                coords *= ANGSTROM_TO_BOHR
            atoms.append(Atom(parts[0], coords))
        return cls(atoms, charge, spin)

    def center(self) -> np.ndarray:
        """Geometric center of the molecule."""
        return np.mean([a.coords for a in self.atoms], axis=0)

    def __repr__(self):
        lines = [f"Molecule(charge={self.charge}, spin={self.spin}, n_elec={self.n_electrons})"]
        for a in self.atoms:
            c = a.coords
            lines.append(f"  {a.symbol:2s}  {c[0]:12.6f} {c[1]:12.6f} {c[2]:12.6f}  (Bohr)")
        return "\n".join(lines)