import numpy as np
from rdkit.Chem import rdPartialCharges, MolFromPDBBlock
from Bio.PDB import Model, Chain, Residue
from Bio.PDB.PDBIO import PDBIO
from io import StringIO

def atoms_to_res(atoms):
    # Convert a list of atom objects into a list of residue objects
    res_dict = {}
    for atom in atoms:
        resname = atom.get_parent().get_resname()
        seg_id = atom.get_parent().get_id()[0]
        res_id = atom.get_parent().get_id()

        # Ensure duplicate residues are not created
        res_key = (res_id, resname, seg_id)
        if res_key not in res_dict.keys():
            # Create new reisdue object
            current_residue = Residue.Residue(res_id, resname, seg_id)
            res_dict[res_key] = current_residue

        else:
            # Retrieve existing residue
            current_residue = res_dict[res_key]

        # Add atom to current residue
        # if atom.get_id() not in [a for a in current_residue.get_atoms()]:
        current_residue.add(atom)
    res_list = list(res_dict.values())
    return res_list

def res_to_model(res_list, chain_id: str = 'A'):
    # Convert a list of residue objects into a model object
    model = Model.Model(id=0)
    chain = Chain.Chain(chain_id)

    seen_res_id = []
    for res in res_list:
        if res.get_id()[1] not in seen_res_id:
            seen_res_id.append(res.get_id()[1])
            chain.add(res)
    model.add(chain)
    return model

def new_pdb_file(model, pdb_id: str = 'A'):
    # Convert model object into a PDB file
    if isinstance(model, Model.Model):
        io = PDBIO()
        io.set_structure(model)
        pdb_file = StringIO()
        io.save(pdb_file)
        pdb_file.seek(0)
        # print(f'Sucessfuly saved your biomolecule as a string-formatted pdb entry.')
    else:
        # print(f'Invalid object of type: {model}.')
        pdb_file = None
    return pdb_file
    
def get_features_charges(coords, atoms):
    # Get partial charges of biomolecule object using RDkit

    res_list = atoms_to_res(atoms)
    model = res_to_model(res_list)
    pdb = new_pdb_file(model)
    mol = MolFromPDBBlock(pdb.getvalue(), sanitize=True)

    num_atoms = mol.GetNumAtoms()
    rdPartialCharges.ComputeGasteigerCharges((mol))

    charges = np.array([float(atom.GetProp('_GasteigerCharge')) for atom in mol.GetAtoms()])
    nan_mask = np.isnan(charges)

    if nan_mask.any():

        charges[nan_mask] = 0.0
    
    return charges

