from collections import namedtuple

hbonds_info = namedtuple('hbonds_info', 'ldon_id lig_don_id prot_acc_id pdon_id prot_don_id lig_acc_id')
hydrophobic_info = namedtuple('hydrophobic_info', 'bs_ids lig_ids pairs_ids')
halogen_info = namedtuple('halogen_info', 'don_id acc_id')
pistack_info = namedtuple('pistack_info', 'proteinring_atoms, proteinring_center ligandring_atoms '
                                          'ligandring_center type')
pication_info = namedtuple('pication_info', 'ring_center charge_center ring_atoms charge_atoms, protcharged')
sbridge_info = namedtuple('sbridge_info', 'positive_atoms negative_atoms positive_center negative_center protispos')
wbridge_info = namedtuple('wbridge_info', 'don_id acc_id water_id protisdon')
metal_info = namedtuple('metal_info', 'metal_id, target_id location')


class VisualizerData:
    """Contains all information on a complex relevant for visualization. Can be pickled"""

    def __init__(self, mol, site):
        pcomp = mol
        pli = mol.interaction_sets[site]
        ligand = pli.ligand

        # General Information
        self.lig_members = sorted(pli.ligand.members)
        self.sourcefile = pcomp.sourcefiles['pdbcomplex']
        self.corrected_pdb = pcomp.corrected_pdb
        self.pdbid = mol.pymol_name
        self.hetid = ligand.hetid
        self.ligandtype = ligand.type
        self.chain = ligand.chain if not ligand.chain == "0" else ""  # #@todo Fix this
        self.position = str(ligand.position)
        self.uid = ":".join([self.hetid, self.chain, self.position])
        self.outpath = mol.output_path
        self.metal_ids = [x.m_orig_idx for x in pli.ligand.metals]
        self.unpaired_hba_idx = pli.unpaired_hba_orig_idx
        self.unpaired_hbd_idx = pli.unpaired_hbd_orig_idx
        self.unpaired_hal_idx = pli.unpaired_hal_orig_idx

        # Information on Interactions

        # Hydrophobic Contacts
        # Contains IDs of contributing binding site, ligand atoms and the pairings
        hydroph_pairs_id = [(h.bsatom_orig_idx, h.ligatom_orig_idx) for h in pli.hydrophobic_contacts]
        self.hydrophobic_contacts = hydrophobic_info(bs_ids=[hp[0] for hp in hydroph_pairs_id],
                                                     lig_ids=[hp[1] for hp in hydroph_pairs_id],
                                                     pairs_ids=hydroph_pairs_id)

        # Hydrogen Bonds
        # #@todo Don't use indices, simplify this code here
        hbonds_ldon, hbonds_pdon = pli.hbonds_ldon, pli.hbonds_pdon
        hbonds_ldon_id = [(hb.a_orig_idx, hb.d_orig_idx) for hb in hbonds_ldon]
        hbonds_pdon_id = [(hb.a_orig_idx, hb.d_orig_idx) for hb in hbonds_pdon]
        self.hbonds = hbonds_info(ldon_id=[(hb.a_orig_idx, hb.d_orig_idx) for hb in hbonds_ldon],
                                  lig_don_id=[hb[1] for hb in hbonds_ldon_id],
                                  prot_acc_id=[hb[0] for hb in hbonds_ldon_id],
                                  pdon_id=[(hb.a_orig_idx, hb.d_orig_idx) for hb in hbonds_pdon],
                                  prot_don_id=[hb[1] for hb in hbonds_pdon_id],
                                  lig_acc_id=[hb[0] for hb in hbonds_pdon_id])

        # Halogen Bonds
        self.halogen_bonds = [halogen_info(don_id=h.don_orig_idx, acc_id=h.acc_orig_idx)
                              for h in pli.halogen_bonds]

        # Pistacking
        self.pistacking = [pistack_info(proteinring_atoms=pistack.proteinring.atoms_orig_idx,
                                        proteinring_center=pistack.proteinring.center,
                                        ligandring_atoms=pistack.ligandring.atoms_orig_idx,
                                        ligandring_center=pistack.ligandring.center,
                                        type=pistack.type) for pistack in pli.pistacking]

        # Pi-cation interactions
        self.pication = [pication_info(ring_center=picat.ring.center,
                                       charge_center=picat.charge.center,
                                       ring_atoms=picat.ring.atoms_orig_idx,
                                       charge_atoms=picat.charge.atoms_orig_idx,
                                       protcharged=picat.protcharged)
                         for picat in pli.pication_paro + pli.pication_laro]

        # Salt Bridges
        self.saltbridges = [sbridge_info(positive_atoms=sbridge.positive.atoms_orig_idx,
                                         negative_atoms=sbridge.negative.atoms_orig_idx,
                                         positive_center=sbridge.positive.center,
                                         negative_center=sbridge.negative.center,
                                         protispos=sbridge.protispos)
                            for sbridge in pli.saltbridge_lneg + pli.saltbridge_pneg]

        # Water Bridgese('wbridge_info', 'don_id acc_id water_id protisdon')
        self.waterbridges = [wbridge_info(don_id=wbridge.d_orig_idx,
                                          acc_id=wbridge.a_orig_idx,
                                          water_id=wbridge.water_orig_idx,
                                          protisdon=wbridge.protisdon) for wbridge in pli.water_bridges]

        # Metal Complexes
        self.metal_complexes = [metal_info(metal_id=metalc.metal_orig_idx,
                                           target_id=metalc.target_orig_idx,
                                           location=metalc.location) for metalc in pli.metal_complexes]