'''For generating periodically-tiled topologies from RDKit Mols'''

__author__ = 'Timotej Bernat'
__email__ = 'timotej.bernat@colorado.edu'

import logging
LOGGER = logging.getLogger(__name__)

import numpy as np
from rdkit.Chem import Mol, CombineMols
from rdkit.Chem.rdMolTransforms import ComputeCanonicalTransform, TransformConformer

from ...genutils.typetools.numpytypes import Shape, N
from ...maths.linearalg import affine


def rdmol_effective_radius(rdmol : Mol, conf_id : int=0) -> float:
    '''Determine an effective radius of influence of a molecule such that all atoms are within this radius'''
    conformer = rdmol.GetConformer(conf_id)
    positions = conformer.GetPositions()
    
    positions_centered = positions - positions.mean(axis=0) # note this is mass-unweighted (i.e. not the center of mass)
    dists_to_centroid = np.linalg.norm(positions_centered, axis=1)
    
    return dists_to_centroid.max()

def tile_lattice_with_rdmol(rdmol : Mol, lattice_points : np.ndarray[Shape[N, 3], float], rotate_randomly : bool=True, conf_id : int=0) -> Mol:
    '''
    Generate a tiled topology of copies of an RDKit molecule, transformed to occupy the same relative positions as points on the given lattice with unit dimensions
    if "random_rotations" is True, each occurrence of the molecule will also have a random rotation applied to it. Mol is NOT modified at any point in the procedure
    '''
    if rdmol.GetNumConformers() < 1:
        raise ValueError('Molecule must have at least one conformer to be tiled')
    
    conformer = rdmol.GetConformer(conf_id)
    centering = ComputeCanonicalTransform(conformer, ignoreHs=True) # translation which centers the given conformer
    
    rotation = centering[:-1, :-1]
    orient_sign = np.sign(np.linalg.det(rotation))
    if orient_sign < 0.0:
        centering[:-1, :-1] *= orient_sign # enforce right-handed coordinate system to avoid unintended inversions
        LOGGER.warning(
            'Caught and corrected stereochemical inversion during RDKit alignment transform;\n' 
            'recommend upgrading RDKit to >=2025.09.x(see https://github.com/rdkit/rdkit/issues/8720)'
        )

    tiled_topology = None
    for point in lattice_points:
        translation = affine.xyzTrans(*point) # translation to the current point on the lattice
        rotation    = affine.randRot(about_x=rotate_randomly, about_y=rotate_randomly, about_z=rotate_randomly) # rotation about either all axes or the identity matrix
        transform = translation @ rotation @ centering # NOTE : (right-to-left) order matters here!! Must first move to origin, then rotate, and then translate to latice point 

        clone = Mol(rdmol) # make copy to avoid modifying ariginal
        TransformConformer(clone.GetConformer(0), transform) # apply transform to 
        
        if tiled_topology is None:
            tiled_topology = clone
        else:
            tiled_topology = CombineMols(tiled_topology, clone)

    return tiled_topology