import os
import random
import numpy as np
from ase import Atom, Atoms
from ase.data import covalent_radii, chemical_symbols
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist

# =============================================================================
# Distance and geometry utilities
# =============================================================================

def distance_atoms2atoms(atoms1, atoms2):
    """
    Minimum normalized interatomic distance between two Atoms objects.
    Distances are scaled by the sum of covalent radii.
    """
    min_distance = float('inf')
    for atom1 in atoms1:
        r1 = covalent_radii[atom1.number]
        for atom2 in atoms2:
            r2 = covalent_radii[atom2.number]
            d = np.linalg.norm(atom1.position - atom2.position) / (r1 + r2)
            if d < min_distance:
                min_distance = d
    return min_distance


def distance_atom2atoms(atom1, atoms2):
    """
    Minimum normalized distance between a single atom and an Atoms object.
    """
    r1 = covalent_radii[atom1.number]
    min_distance = float('inf')
    for atom2 in atoms2:
        r2 = covalent_radii[atom2.number]
        d = np.linalg.norm(atom1.position - atom2.position) / (r1 + r2)
        if d < min_distance:
            min_distance = d
    return min_distance


def distance_interatom(atoms):
    """
    Minimum normalized distance between any pair of atoms in the same molecule.
    """
    min_distance = float('inf')
    n = len(atoms)
    for i in range(n):
        r1 = covalent_radii[atoms[i].number]
        for j in range(i + 1, n):
            r2 = covalent_radii[atoms[j].number]
            d = np.linalg.norm(atoms[i].position - atoms[j].position) / (r1 + r2)
            if d < min_distance:
                min_distance = d
    return min_distance


def centroid(atoms):
    """Geometric centroid (not center of mass)."""
    return np.mean(atoms.get_positions(), axis=0)


def radius_max(atoms):
    """
    Maximum distance from centroid including covalent radius.
    Useful as an effective cluster size.
    """
    ctd = centroid(atoms)
    radii = [
        np.linalg.norm(atom.position - ctd) + covalent_radii[atom.number]
        for atom in atoms
    ]
    return max(radii)


def scale_coords(atoms, factor):
    """
    Uniformly scale coordinates around the center of mass.
    """
    atoms_out = atoms.copy()
    rel = atoms.positions - atoms.get_center_of_mass()
    atoms_out.positions = rel * factor
    return atoms_out


# =============================================================================
# Graph and connectivity utilities
# =============================================================================

def adjacency_matrix(atoms, factor=1.2):
    """
    Binary adjacency matrix based on scaled covalent distances.
    """
    n = len(atoms)
    mat = np.zeros((n, n), dtype=np.int64)

    for i in range(n):
        ri = covalent_radii[atoms[i].number]
        pi = atoms[i].position
        for j in range(i + 1, n):
            rj = covalent_radii[atoms[j].number]
            pj = atoms[j].position
            d = np.linalg.norm(pj - pi) / (ri + rj)
            if d <= factor:
                mat[i, j] = mat[j, i] = 1
    return mat


def degree_matrix(atoms, factor=1.2):
    """
    Degree matrix associated with the adjacency matrix.
    """
    adj = adjacency_matrix(atoms, factor)
    return np.diag(np.sum(adj, axis=1))


def connected_graph(atoms, factor=1.2):
    """
    Checks if the molecular graph is fully connected using
    repeated adjacency propagation.
    """
    adj = adjacency_matrix(atoms, factor)
    n = len(adj)

    visited = np.zeros(n, dtype=int)
    visited[0] = 1
    prev = visited.copy()

    while True:
        visited = np.dot(adj, visited) + prev
        visited = np.clip(visited, 0, 1)
        if np.array_equal(visited, prev):
            break
        prev = visited

    return int(np.sum(visited) == n)


# =============================================================================
# Rotations and random orientations
# =============================================================================

def rand_unit_vector():
    """Random unit vector uniformly distributed on the sphere."""
    phi = random.uniform(0.0, 2.0 * np.pi)
    theta = random.uniform(0.0, np.pi)
    return np.array([
        np.sin(theta) * np.cos(phi),
        np.sin(theta) * np.sin(phi),
        np.cos(theta)
    ])


def random_deg_angles():
    """Random Euler angles in degrees."""
    return (
        random.randint(0, 360),
        random.randint(0, 360),
        random.randint(0, 360)
    )


def euler_matrix(qdegx, qdegy, qdegz):
    """Rotation matrix from Euler angles (degrees)."""
    rx, ry, rz = np.deg2rad([qdegx, qdegy, qdegz])
    cx, cy, cz = np.cos([rx, ry, rz])
    sx, sy, sz = np.sin([rx, ry, rz])

    return np.array([
        [cy * cz, -cy * sz, sy],
        [cx * sz + cz * sx * sy, cx * cz - sx * sy * sz, -cy * sx],
        [sx * sz - cx * cz * sy, cz * sx + cx * sy * sz, cx * cy]
    ])


def rotate_matrix(molecule, matrix):
    """Apply a rotation matrix to atomic positions."""
    molecule.set_positions(np.dot(molecule.get_positions(), matrix.T))
    return molecule


def rotate_deg(molecule, qdegx, qdegy, qdegz):
    """Rotate molecule using Euler angles."""
    return rotate_matrix(molecule, euler_matrix(qdegx, qdegy, qdegz))


def rotate_random(molecule):
    """Apply a random rotation."""
    return rotate_deg(molecule, *random_deg_angles())


def rodrigues_rotation_matrix(kvector, qdeg):
    """Rodrigues rotation matrix around arbitrary axis."""
    qrad = np.deg2rad(qdeg)
    k = np.array(kvector, dtype=float)
    k /= np.linalg.norm(k)

    K = np.array([
        [0, -k[2], k[1]],
        [k[2], 0, -k[0]],
        [-k[1], k[0], 0]
    ])

    return np.eye(3) + np.sin(qrad) * K + (1 - np.cos(qrad)) * np.dot(K, K)


def rotate_vector_angle_deg(molecule, kvector, qdeg):
    """Rotate molecule around axis kvector by qdeg."""
    return rotate_matrix(molecule, rodrigues_rotation_matrix(kvector, qdeg))


# =============================================================================
# Alignment and RMSD
# =============================================================================

def align(atoms):
    """
    Align molecule along principal inertia axes.
    """
    atoms = atoms.copy()
    atoms.translate(-centroid(atoms))
    _, vecs = atoms.get_moments_of_inertia(vectors=True)
    atoms.set_positions(np.dot(atoms.get_positions(), vecs.T))
    return atoms


def KuhnMunkres(mol1, mol2, use_elements=True):
    """
    Atom matching using Hungarian algorithm.
    Large penalties enforce element-wise matching.
    """
    p1 = mol1.positions - mol1.get_center_of_mass()
    p2 = mol2.positions - mol2.get_center_of_mass()

    cost = cdist(p1, p2)

    if use_elements:
        s1 = mol1.get_chemical_symbols()
        s2 = mol2.get_chemical_symbols()
        for i in range(len(s1)):
            for j in range(len(s2)):
                if s1[i] != s2[j]:
                    cost[i, j] = 1e10

    _, col = linear_sum_assignment(cost)

    mol2a = mol2.copy()
    mol2a.symbols = [mol2.symbols[i] for i in col]
    mol2a.set_positions(mol2.positions[col])
    return mol2a


def rmsd(mol1, mol2):
    """Root-mean-square deviation after centering."""
    p1 = mol1.positions - mol1.get_center_of_mass()
    p2 = mol2.positions - mol2.get_center_of_mass()
    return np.sqrt(np.mean(np.sum((p1 - p2) ** 2, axis=1)))


# NOTE:
# align_two is intentionally left algorithmically unchanged.
# It explores Cartesian inversion symmetries + atom permutation
# to ensure a global minimum RMSD.

def align_two(atoms1, atoms2):
    mol1 = align(atoms1.copy())
    mol2 = align(atoms2.copy())
    sym_mats = [
        np.diag([ sx, sy, sz ])
        for sx in (1, -1)
        for sy in (1, -1)
        for sz in (1, -1)
    ]
    min_rms = float('inf')
    best = None
    for M in sym_mats:
        mol2x = mol2.copy()
        mol2x.set_positions(np.dot(mol2x.positions, M.T))
        mol2a = KuhnMunkres(mol1, mol2x)
        r = rmsd(mol1, mol2a)
        if r < min_rms:
            min_rms = r
            best = mol2a.copy()
    return mol1, best


def align_list(atomlist):
    """Sequential alignment of a list of molecules."""
    if len(atomlist) == 0:
        return []
    out = [align(atomlist[0])]
    for i in range(1, len(atomlist)):
        _, mol = align_two(out[-1], atomlist[i])
        out.append(mol)
    return out

def planarity_index(atoms):
    """
    Ratio indicating molecular planarity based on inertia tensor.
    """
    ia, ib, ic = atoms.get_moments_of_inertia()
    return ic / (ia + ib)

# =============================================================================
# IO utilities (XYZ)
# =============================================================================

def readxyzs(filename):
    """
    Read concatenated XYZ file into a list of Atoms objects.
    Energy and label are stored in atoms.info.
    """
    if not os.path.isfile(filename):
        raise FileNotFoundError(filename)
    molecules = []
    with open(filename) as fh:
        while True:
            line = fh.readline()
            if not line:
                break
            natoms = int(line.strip())
            comment = fh.readline().split()
            name = comment[1] if len(comment) >= 2 else "unknown"
            energy = float(comment[0]) if comment else 0.0
            mol = Atoms()
            mol.info["e"] = energy
            mol.info["i"] = name
            for _ in range(natoms):
                sym_raw, x, y, z = fh.readline().split()
                x, y, z = map(float, (x, y, z))
                try:
                    symbol = chemical_symbols[int(sym_raw)]
                except (ValueError, IndexError):
                    symbol = sym_raw.capitalize()
                mol.append(Atom(symbol, (x, y, z)))
            molecules.append(mol)
    return molecules

def writexyzs(atoms_list, filename):
    """Write list of Atoms objects to concatenated XYZ file."""
    if not isinstance(atoms_list, list):
        atoms_list = [atoms_list]
    with open(filename, "w") as fh:
        for atoms in atoms_list:
            print(len(atoms), file=fh)
            print(f"{atoms.info['e']:12.8f} {atoms.info['i']}", file=fh)
            for atom in atoms:
                x, y, z = atom.position
                print(f"{atom.symbol:<2s} {x:16.9f} {y:16.9f} {z:16.9f}", file=fh)

# =============================================================================
# List and filtering utilities
# =============================================================================

def rename(atoms_list, basename, ndigits):
    """Rename molecules with zero-padded indices."""
    for i, atoms in enumerate(atoms_list):
        atoms.info["i"] = f"{basename}_{str(i+1).zfill(ndigits)}"
    return atoms_list

def sort_by_energy(atoms_list, opt=0):
    """
    Sort molecules by energy.
    Optionally shift energies relative to minimum.
    """
    if not atoms_list: return []
    energies = [(i, a.info["e"]) for i, a in enumerate(atoms_list)]
    energies.sort(key=lambda x: x[1])
    ref = energies[0][1] if opt == 0 else 0.0
    out = []
    for idx, e in energies:
        a = atoms_list[idx].copy()
        a.info["e"] = e - ref
        out.append(a)
    return out

def cutter_nonconnected(atoms_list, factor=1.2):
    """Remove disconnected structures."""
    return [a for a in atoms_list if connected_graph(a, factor)]

def cutter_energy(atoms_list, enemax):
    """Filter structures within energy window."""
    sorted_list = sort_by_energy(atoms_list, 1)
    emin = sorted_list[0].info["e"]
    return [a for a in atoms_list if (a.info["e"] - emin) < enemax]

def merge_atoms(atoms_list):
    atomsout = atoms_list[0].copy()
    for index in range(1, len(atoms_list)):
        for iatom in atoms_list[index]:
            atomsout.append(iatom)
    return atomsout

# =============================================================================
# Parallelization helpers
# =============================================================================

def listflatten(ntotal, nproc):
    """Distribute ntotal items over nproc workers."""
    if ntotal <= nproc:
        return [1] * ntotal
    return [
        ntotal // nproc + int((ntotal % nproc) > i)
        for i in range(nproc)
    ]

def prepare_folders(poscarlist, nproc, base_name):
    """Create folders for parallel execution."""
    folderlist = []
    for i in range(nproc):
        name = f"{base_name}proc{str(i+1).zfill(3)}"
        folderlist.append(name)
        os.makedirs(name, exist_ok=True)
    return folderlist

def split_poscarlist(poscarlist, nproc):
    """Split list into nproc chunks."""
    sizes = listflatten(len(poscarlist), nproc)
    out, idx = [], 0
    for s in sizes:
        out.append(poscarlist[idx:idx+s])
        idx += s
    return out
