from rdkit.Geometry import Point3D
from rdkit.Chem.rdmolfiles import PDBWriter
from rdkit import Chem
import numpy as np
from functools import reduce





class Molecule:

    def __init__(self, mol, coordinates, radii, is_salt=False, is_transition_metal_complex=False):

        self.mol = mol
        self.coordinates = coordinates
        self.set_mol_conf_coordinates(coordinates)
        self.radii = radii
        self.is_salt = is_salt

    def set_radii(self, radii):
        self.radii = radii

    def set_coordinates(self, coordinates):

        assert np.array(self.coordinates).shape == np.array(coordinates).shape, (
            str(np.array(self.coordinates).shape)
            + " does not match"
            + str(np.array(coordinates).shape)
        )
        self.coordinates = coordinates

    def set_mol_conf_coordinates(self, coordinates):
        """Sets coordinates for the Molecule and RDMol object

        Args:
            coordinates (_type_): list with coordinates

        Returns:
            _type_: Rdkit Conformer
        """

        self.set_coordinates(coordinates=coordinates)
        conf = self.mol.GetConformer()
        for i in range(self.mol.GetNumAtoms()):
            x, y, z = coordinates[i]
            conf.SetAtomPosition(i, Point3D(x, y, z))
        return conf

    def get_atom_symbols(self):
        
        symbols = [atom.GetSymbol() for atom in self.mol.GetAtoms()]
        return symbols

    def get_atom_numbers(self): 
        
        numbers = [atom.GetAtomicNum() for atom in self.mol.GetAtoms()]
        return numbers 

    def is_salt(self): 
        return self.is_salt


class TransitionMetalComplex(Molecule): 

    def __init__(self, mol, coordinates, radii): 
        self.mol = mol
        self.coordinates = coordinates
        self.set_mol_conf_coordinates(coordinates)
        self.radii = radii
        self.is_salt = is_salt
        self.is_transition_metal_complex = True



class MoleculeCollection:

    def __init__(self, molecules, is_salt=False):
        self.molecules = molecules
        self.coordinates = self.get_atom_coords_list()
        self.radii = self.get_radii_list()
        self.atom_symbols = self.get_atom_symbols()
        self.is_salt = is_salt

    def add_molecule(self, molecule) -> None:

        self.molecules.append(molecule)

    def get_atom_coords_list(self) -> list:

        # also kind of flattened such that it is n_atoms, 3 in shape
        atom_list = []
        for molecule in self.molecules:
            coords = molecule.coordinates
            atom_list += coords
        return atom_list

    def get_atom_symbols(self) -> list:

        atom_symbols = []
        for molecule in self.molecules:
            atom_symbols += molecule.get_atom_symbols()
        return atom_symbols

    def get_atom_numbers(self) -> list:

        atom_numbers = []
        for molecule in self.molecules:
            atom_numbers += molecule.get_atom_numbers()
        return atom_numbers


    def get_radii_list(self) -> list:

        radii_list = []
        for molecule in self.molecules:
            radii_list += molecule.radii
        return radii_list

    def get_mols_list(self) -> list:
        mols_list = []
        for molecule in self.molecules:
            mols_list.append(molecule.mol)
        return mols_list

    def join_rd_mols(self):

        mols_list = self.get_mols_list()
        mol = reduce(Chem.CombineMols, mols_list)
        return mol

    def save_to_xyz(self, file_name: str) -> None:
        atom_symbols = self.get_atom_symbols()
        coordinates = self.get_atom_coords_list()
        with open(file_name, "w") as f:
            f.write(f"{len(atom_symbols)}\n")
            f.write("XYZ file generated by mond\n")
            for symbol, (x, y, z) in zip(atom_symbols, coordinates):
                f.write(f"{symbol} {x:.6f} {y:.6f} {z:.6f}\n")

    def save_to_pdb(self, file_name: str) -> None:

        mols = self.get_mols_list()

        with PDBWriter(file_name) as writer:
            for mol in mols:
                writer.write(mol)

    def save_single_mol_to_pdb(self, file_name: str) -> None:

        mol = self.join_rd_mols()

        with PDBWriter(fileName=file_name) as writer:
            writer.write(mol)

    def save_to_sdf(self, file_name: str) -> None:

        mols = self.get_mols_list()

        with Chem.SDWriter(file_name) as writer:
            for mol in mols:
                writer.write(mol)

    def save_single_mol_to_sdf(self, file_name: str) -> None:

        mol = self.join_rd_mols()

        with Chem.SDWriter(file_name) as writer:
            writer.write(mol)


class SaltCollection(MoleculeCollection): 
    
    def __init__(self, anions:list, cations:list):

        self.anions = anions
        self.cations = cations
        self.molecules = anions + cations
        self.coordinates = self.get_atom_coords_list()
        self.radii = self.get_radii_list()
        self.atom_symbols = self.get_atom_symbols()
        self.is_salt = True