"""Crystal structure extraction, building, and supercell expansion."""

import re
import numpy as np
from typing import Dict, List, Optional

from pymatgen.core import Structure
from pymatgen.core.periodic_table import Specie

from ceraviz.types import (
    AtomTD, LatticeTD, CrystalTD, MaterialMetaTD, MaterialTD,
)


def parse_oxidation_states(possible_species: Optional[List[str]]) -> Dict[str, int]:
    """Parse oxidation state strings like 'Fe3+' or 'O2-' into a dict."""
    oxi: Dict[str, int] = {}
    if not possible_species:
        return oxi

    for s in possible_species:
        s = s.replace(" ", "")
        m = re.match(r"^([A-Za-z]{1,2})(\d*)([+-])$", s)
        if m:
            elem, mag, sign = m.groups()
            mag = int(mag) if mag else 1
            oxi[elem] = mag if sign == "+" else -mag
            continue

        m = re.match(r"^([A-Za-z]{1,2})([+-]\d+)$", s)
        if m:
            oxi[m.group(1)] = int(m.group(2))

    return oxi


def extract_materials(docs) -> List[MaterialTD]:
    """Extract typed material data from Materials Project API documents."""
    materials: List[MaterialTD] = []

    for doc in docs:
        s = doc.structure

        lattice: LatticeTD = {
            "abc": tuple(s.lattice.abc),
            "angles": tuple(s.lattice.angles),
            "vectors": np.array(s.lattice.matrix),
            "volume": s.lattice.volume,
            "pbc": tuple(s.lattice.pbc),
        }

        atoms: List[AtomTD] = [
            {
                "element": site.specie.symbol,
                "frac": np.array(site.frac_coords),
                "cart": np.array(site.coords),
            }
            for site in s.sites
        ]

        structure: CrystalTD = {
            "lattice": lattice,
            "atoms": atoms,
            "nsites": doc.nsites,
            "elements": [el.symbol for el in doc.elements],
        }

        meta: MaterialMetaTD = {
            "material_id": str(doc.material_id),
            "formula": doc.formula_pretty,
            "chemsys": doc.chemsys,
            "symmetry_symbol": doc.symmetry.symbol,
            "symmetry_number": doc.symmetry.number,
            "crystal_system": doc.symmetry.crystal_system.value,
            "density": doc.density,
            "is_stable": doc.is_stable,
            "energy_above_hull": doc.energy_above_hull,
        }

        materials.append({
            "structure": structure,
            "meta": meta,
        })

    return materials


def make_supercell(structure_td: CrystalTD, nx: int = 2, ny: int = 2, nz: int = 2) -> CrystalTD:
    """Expand structure to nx×ny×nz supercell."""
    lattice = structure_td["lattice"]
    old_vectors = np.array(lattice["vectors"])

    new_vectors = old_vectors.copy()
    new_vectors[0] *= nx
    new_vectors[1] *= ny
    new_vectors[2] *= nz

    new_abc = (
        lattice["abc"][0] * nx,
        lattice["abc"][1] * ny,
        lattice["abc"][2] * nz,
    )
    new_volume = lattice["volume"] * nx * ny * nz

    new_atoms: List[AtomTD] = []
    for ix in range(nx):
        for iy in range(ny):
            for iz in range(nz):
                shift = ix * old_vectors[0] + iy * old_vectors[1] + iz * old_vectors[2]
                for atom in structure_td["atoms"]:
                    new_frac = (np.array(atom["frac"]) + np.array([ix, iy, iz])) / np.array([nx, ny, nz])
                    new_atoms.append({
                        "element": atom["element"],
                        "frac": new_frac,
                        "cart": np.array(atom["cart"]) + shift,
                    })

    new_lattice: LatticeTD = {
        "abc": new_abc,
        "angles": lattice["angles"],
        "vectors": new_vectors,
        "volume": new_volume,
        "pbc": lattice["pbc"],
    }

    return {
        "lattice": new_lattice,
        "atoms": new_atoms,
        "nsites": len(new_atoms),
        "elements": structure_td["elements"],
    }


def build_pmg_structure(
    structure_td: CrystalTD,
    oxidation_states: Optional[Dict[str, int]] = None,
) -> Structure:
    """Build a pymatgen Structure from typed dict data."""
    lattice = np.array(structure_td["lattice"]["vectors"])

    species = [
        Specie(a["element"], oxidation_states[a["element"]])
        if oxidation_states and a["element"] in oxidation_states
        else a["element"]
        for a in structure_td["atoms"]
    ]

    frac_coords = [a["frac"] for a in structure_td["atoms"]]

    return Structure(
        lattice=lattice,
        species=species,
        coords=frac_coords,
        coords_are_cartesian=False,
        to_unit_cell=True,
    )
