# This is a separate module because building ira_mod [1] needs a Makefile and a
# PYTHONPATH export.. after that's cleared up this will probably be a dependency
# [1]: https://github.com/mammasmias/IterativeRotationsAssignments
import dataclasses
from collections import Counter

import ira_mod
import numpy as np


@dataclasses.dataclass
class IRAComp:
    """Result of an IRA comparison (rotation, translation, permutation, Hausdorff distance).

    .. versionadded:: 1.0.0
    """

    rot: np.array
    trans: np.array
    perm: np.array
    hd: float


class IncomparableStructuresError(ValueError):
    """Custom exception raised for incompatible atomistic structures.

    .. versionadded:: 1.0.0
    """

    pass


def _perform_ira_match(atm1, atm2, k_factor=2.8):
    """Performs IRA matching. Internal function."""
    if len(atm1) != len(atm2):
        errmsg = "Atomistic structures must have the same number of atoms."
        raise IncomparableStructuresError(errmsg)
    if not Counter(atm1.symbols) == Counter(atm2.symbols):
        errmsg = "Atomistic structures must have the same atom types."
        raise IncomparableStructuresError(errmsg)

    ira = ira_mod.IRA()
    return ira.match(
        len(atm1),
        atm1.get_atomic_numbers(),
        atm1.get_positions(),
        len(atm2),
        atm2.get_atomic_numbers(),
        atm2.get_positions(),
        k_factor,
    )


def is_ira_pair(atm1, atm2, hd_tol=1, k_factor=2.8):
    """Checks if two atomistic structures are an IRA pair.

    .. versionadded:: 1.0.0
    """
    try:
        _, _, _, hd = _perform_ira_match(atm1, atm2, k_factor)
        return hd < hd_tol
    except IncomparableStructuresError:
        return False


def do_ira(atm1, atm2, k_factor=2.8):
    """Performs IRA matching on two atomistic structures.

    .. versionadded:: 1.0.0
    """
    rotation, translation, perm, hd = _perform_ira_match(atm1, atm2, k_factor)
    return IRAComp(rot=rotation, trans=translation, perm=perm, hd=hd)


def calculate_rmsd(atm1, atm2, k_factor=2.8):
    """Calculates RMSD using do_ira.

    .. versionadded:: 1.0.0
    """
    ira_comp = do_ira(atm1, atm2, k_factor)
    atm2_coords_permuted = atm2.get_positions()[ira_comp.perm]
    atm2_coords_transformed = np.dot(atm2_coords_permuted, ira_comp.rot) + ira_comp.trans
    n_atoms = len(atm1.get_positions())
    squared_distances = np.sum(
        (atm1.get_positions() - atm2_coords_transformed) ** 2, axis=1
    )
    rmsd = np.sqrt(np.sum(squared_distances) / n_atoms)
    return rmsd
