#!/usr/bin/env python
import numpy as np
from openbabel import pybel
# import time
from rdkit import Chem
from rdkit.Chem import Draw
from pdbtools.pdbtools import pdbtools

def cal_angle_from_points(dist_lp, dist_dh, dist_ah):
    cos_theta = (np.power(dist_dh, 2) + np.power(dist_ah, 2) -
                 np.power(dist_lp, 2)) / (2 * dist_dh * dist_ah)
    theta = np.arccos(cos_theta)

    return theta

def cal_angle_from_vectors(A, B):
    cos_theta = np.dot(A, B)
    theta = np.arccos(cos_theta)
    return theta

class Pharmacophore(object):
    """
    pharmacophore
    """
    feature_type_dict = {
        'HBD': 0,
        'HBA': 1,
        'Anion': 2,
        'Cation': 3,
        'Hydrophobic': 4,
        'Aromatic': 5
    }
    feature_type_list = ['HBD', 'HBA', 'Anion',
                         'Cation', 'Hydrophobic', 'Aromatic']

    complementary_feature = {
        'HBD': [['HBA', 3.4]],
        'HBA': [['HBD', 3.4]],
        'Anion': [['Cation', 5.0]],
        'Cation': [['Anion', 5.0], ['Aromatic', 5.0]],
        'Hydrophobic': [['Hydrophobic', 6.5]],
        'Aromatic': [['Cation', 5.0], ['Aromatic', 6.0]],
        'MBA': [['Metal', 3.1]],
        'Metal': [['MBA', 3.1]],
    }

    num_feature_type = len(feature_type_list)
    SMARTS_pattern_dict = {
        'HBD': {
            0: '[N&H2&v3]',
            1: '[N&H1&v3]',
            2: '[N&!H0&+1&v4]',
            3: '[n&H1&+0]',
            4: '[n&H1&+1]',
            5: '[O;H1;+0]',
            6: '[S;H1;+0]',
            7: '[O;H2;+0]',
        },
        'HBA': {
            0: '[$([O,S;-])]',
            1: '[$([O,S;H0;v2])]',
            2: '[$([O,S;H1;v2]-[!$(*=[O,N,P,S])])]',
            3: '[$([N;v3;!$(N-*=!@[O,N,P,S])])]',
            4: '[$([nH0,o,s;+0])]',
            5: '[$([N;H0]#[C&v4])]',
            6: '[N&v3;H0;$(Nc)]',
            7: '[F;$(F-[#6])]',
            8: '[O;H2;+0]',

        },
        'Anion': {
            0: '[-]',
            1: '[SX4](=O)(=O)(-[O;H1,H0&-1])',
            2: '[PX4](=O)(-[O;H1,H0&-1])([!O])([!O])',
            3: '[PX4](=O)(-[O;H1,H0&-1])(-[O;H1,H0&-1])',
            4: '[CX3,SX3](=[O,S,P])-[O;H1,H0&-1]',
        },
        'Cation': {
            0: '[+]',
            1: '[NX3]=[CX3]([NX3])[!N]',
            2: 'NC(=N)N',
            3: 'c1ncnc1'
        },
        'Hydrophobic': {
            0: '[D3,D4;#6;+0;!$([#6][#7,#8,#9])]',
            1: '[R0;D2;#6;+0;!$([#6][#7,#8,#9])]',
            2: '[CX4](F)(F)(F)',
            3: '[#6;R]',
            4: '[#17,#35,#53]',
        },
        'Aromatic': {
            0: 'a1:a:a:a:1',
            1: 'a1:a:a:a:a:1',
            2: 'a1:a:a:a:a:a:1',
            3: 'a1:a:a:a:a:a:a:1',
            4: 'a1:a:a:a:a:a:a:a:1',
        }
    }

    metal_name_dict = {
        'MG': 0,
        'K': 1,
        'MN': 2,
        'FE': 3,
        'ZN': 4,
    }
    metal_type_list = ['MG', 'K', 'MN', 'FE', 'ZN']
    metal_atomic_name_dict = {12: 'MG', 19: 'K', 25: 'MN', 26: 'FE', 30: 'ZN'}
    metal_bind_name_dict = {
        7: 0,
        8: 1,
        16: 2,
    }
    metal_interaction_cutoff = {
        0: {2: 2.6, 3: 2.5, 4: 2.56},
        1: {0: 3.1, 1: 2.4, 2: 2.7, 3: 2.5, 4: 2.8},
        2: {4: 2.6},
    }

    weight_pis = {
        'HBD:HBA': 0.129,
        'HBA:HBD': 0.129,
        'Anion:Cation': 0.31,
        'Cation:Anion': 0.31,
        'Cation:Aromatic': 0.039,
        'Hydrophobic:Hydrophobic': 0.00585,
        'Aromatic:Cation': 0.039,
        'Aromatic:Aromatic': 0.064,
        'MBA:Metal': 1.0
    }
    bias_pis = 4.0
    cutoff = 6.5
    use_pinfo = False

    def __init__(self, params=None):
        """
            initialize Pharmacophore
            set parameters
        """

        self.interaction_feature = list(self.weight_pis.keys())
        if 'weight_pis' in params:
            self.weight_pis = params['weight_pis']
        if 'bias_pis' in params:
            self.bias_pis = params['bias_pis']
        if 'cutoff' in params:
            self.cutoff = params['cutoff']
        cmin = -99999.999
        cmax = 99999.999

        use_box = False
        if 'dock_config' in params:
            box_center, box_size = self.read_dock_config(params['dock_config'])
            box_min = box_center - box_size/2
            box_max = box_center + box_size/2
            cmin = box_min - self.cutoff
            cmax = box_max + self.cutoff
            use_box = True

        if 'ligand_file' in params:
            ligand_file = params['ligand_file']
            ligand_model_dict = pdbtools.read_coor_pdb(
                ligand_file, exclude_Hs=True)
            ligand_dict = ligand_model_dict[1]
            ligand_coor = list(ligand_dict.values())[0]
            cmin, cmax = pdbtools.cal_ligand_size(ligand_coor)
            cmin = cmin - self.cutoff
            cmax = cmax + self.cutoff
            use_box = True

        if 'include_hydrophobic' in params:
            include_hydrophobic = params['include_hydrophobic']

        piscore_receptor = params['piscore_receptor']
        pf_receptor = params['pf_receptor']

        if piscore_receptor is not None:
            ms = pybel.readfile('pdb', piscore_receptor)
            m_protein = list(ms)[0]
            receptor_bond_dict = self.get_bond_info(m_protein)
            PF_dict_protein = self.find_PF(m_protein, receptor_bond_dict,
                                           is_protein=True)
            PF_coor_protein = self.find_PF_coor(m_protein, PF_dict_protein,
                                receptor_bond_dict, is_protein=True)
            metal_coor = self.find_metal_ion(m_protein, is_protein=True)
            PF_coor_protein['Metal'] = metal_coor
            if use_box:
                self.PF_coor_protein = self.box_protein(
                    PF_coor_protein, cmin, cmax)
            else:
                self.PF_coor_protein = PF_coor_protein
            if pf_receptor is not None:
                PF_coor_protein_dict = {0: self.PF_coor_protein}
                self.write_PF(PF_coor_protein_dict,
                              pf_receptor, is_protein=True)
        elif pf_receptor is not None:
            self.PF_coor_protein = self.read_PF(pf_receptor,
                                                is_protein=True)[0]

        pinfo_ligand = params['pinfo_ligand']
        pf_receptor_info = params['pf_receptor_info']
        if pinfo_ligand is not None:
            self.use_pinfo = True
            template_dict = self.find_template(self.PF_coor_protein, pinfo_ligand,
                                               include_hydrophobic=include_hydrophobic)
            self.PF_coor_info = self.select_template_protein(self.PF_coor_protein,
                                                             template_dict, cut_num=1)
            if pf_receptor_info is not None:
                PF_coor_info_dict = {0: self.PF_coor_info}
                self.write_PF(PF_coor_info_dict, pf_receptor_info,
                              is_protein=True)
        elif pf_receptor_info is not None:
            self.use_pinfo = True
            self.PF_coor_info = self.read_PF(
                pf_receptor_info, is_protein=True)[0]

    def get_bond_info(self, m):
        """
            find neighbor atoms in molecule
            input:
                m: pybel molecule object
            output:
                bond_dict: neighbor atom dictionary
        """
        bond_dict = dict()
        mol = m.OBMol

        for i in range(mol.NumBonds()):
            bb = mol.GetBondById(i)
            if bb is None:
                continue
            begin = bb.GetBeginAtomIdx()
            end = bb.GetEndAtomIdx()
            if begin not in bond_dict:
                bond_dict[begin] = [end]
            else:
                bond_dict[begin] += [end]
            if end not in bond_dict:
                bond_dict[end] = [begin]
            else:
                bond_dict[end] += [begin]
        return bond_dict

    def find_PF(self, m, bond_dict, is_protein=False):
        """
            Find pharmacophoric features
            input:
                m: pybel molecule object
                is_protein: protein or not
            output: PF dictionary
        """
        PF_dict = dict()
        atoms = m.atoms
#        bond_dict = self.get_bond_info(m)

        # search pharmacophoric feature using SMARTS pattern
        for feature_type in self.feature_type_list:
            patterns = self.SMARTS_pattern_dict[feature_type]
            PF_dict[feature_type] = list()
            pattern_keys = sorted(patterns.keys())
            for pattern_idx in pattern_keys:
                smarts_pattern = patterns[pattern_idx]
                smarts = pybel.Smarts(smarts_pattern)
                atom_idx_group_list = smarts.findall(m)
                Nfeatures = len(atom_idx_group_list)
                for k in range(Nfeatures):
                    atom_idx_group = atom_idx_group_list[k]

                    fea_dict = dict()
                    fea_dict['atom_idx_group'] = atom_idx_group
                    fea_dict['pattern_idx'] = pattern_idx
                    if feature_type == 'HBD':
                        h_atoms_idx = list()
                        atom_idx = atom_idx_group[0]
                        neighbor_atoms_idx = bond_dict[atom_idx]
                        for neighbor_atom_idx in neighbor_atoms_idx:
                            neighbor_atom = atoms[neighbor_atom_idx - 1]
                            if neighbor_atom.atomicnum != 1:
                                continue
                            h_atoms_idx += [neighbor_atom_idx]
                        if len(h_atoms_idx) == 0:
                            continue
                        fea_dict['h_atoms_idx'] = h_atoms_idx

                    atom_idx = atom_idx_group[0]
                    atom = atoms[atom_idx-1]
                    fea_dict['atom_type'] = atom.type
                    if is_protein:
                        residue = atom.residue
                        chain_id = residue.OBResidue.GetChain()
                        fea_dict['chain_id'] = chain_id
                        residue_name = residue.name
                        fea_dict['residue_name'] = residue_name
                        residue_num = residue.OBResidue.GetNum()
                        fea_dict['residue_num'] = residue_num
                        fea_dict['weight'] = 1.0

                    PF_dict[feature_type] += [fea_dict]

        return PF_dict

    def find_PF_coor(self, m, PF_dict, bond_dict, is_protein=False):
        """
            Find pharmacophoric features
            input:
                m: pybel molecule object
                PF_dict: pharmacophoric feature dictionary
                is_protein: True or False
            output:
                PF_coor: dictionary
        """

        atoms = m.atoms
        PF_coor = dict()
        for feature_type in self.feature_type_list:
            if feature_type not in PF_dict:
                continue
            fea_dict_list = PF_dict[feature_type]
            PF_coor[feature_type] = list()
            Nfeatures = len(fea_dict_list)
            p00_atoms_old = list()
            num_member = list()
            for k in range(Nfeatures):
                fea_dict = fea_dict_list[k]
                num_member += [len(fea_dict['atom_idx_group'])]
            num_member = -np.array(num_member)
            index = num_member.argsort()
            for k in index:
                fea_dict = fea_dict_list[k]
                atom_idx_group = fea_dict['atom_idx_group']
                pattern_idx = fea_dict['pattern_idx']

                intersection = set(atom_idx_group).intersection(p00_atoms_old)
                p00_atoms_old += atom_idx_group
                if len(intersection) > 0 and feature_type != 'Aromatic':
                    continue

                pseudo_atom = []
                for atom_index in atom_idx_group:
                    coor = np.array(atoms[atom_index - 1].coords)
                    pseudo_atom += [coor]
                pseudo_atom = np.array(pseudo_atom)
                pseudo_atom_coor = pseudo_atom.mean(axis=0)

                fea_dict_new = dict()
                fea_dict_new['atom_idx_group'] = atom_idx_group
                fea_dict_new['pattern_idx'] = pattern_idx
                fea_dict_new['pseudo_atom_coor'] = pseudo_atom_coor

                if feature_type == 'HBD':
                    h_atoms = list()
                    atom_idx = atom_idx_group[0]
                    neighbor_heavy_atom_list = list()
                    neighbor_atoms_idx = bond_dict[atom_idx]
                    neighbor_atom_list = list()
                    for neighbor_atom_idx in neighbor_atoms_idx:
                        neighbor_atom = atoms[neighbor_atom_idx - 1]
                        if neighbor_atom.atomicnum == 1:
                            continue
                        nei_coor = np.array(atoms[neighbor_atom_idx - 1].coords)
                        neighbor_atom_list += [[neighbor_atom_idx, nei_coor]]

                    h_atoms_idx = fea_dict['h_atoms_idx']
                    num_neighbor_atom = len(neighbor_atom_list)
                    for h_atom_idx in h_atoms_idx:
                        h_coor = np.array(atoms[h_atom_idx - 1].coords)
                        h_atoms += [[h_atom_idx, h_coor]]
                        #if len(neighbor_atom_list)==1:
                            #nei_coor = neighbor_atom_list[0][1]
                            #vec_a = pseudo_atom_coor - nei_coor
                            #vec_b = h_coor - pseudo_atom_coor
                            #vec_c = h_coor - nei_coor
                            #a = np.linalg.norm(vec_a)
                            #b = np.linalg.norm(vec_b)
                            #c = np.linalg.norm(vec_c)
                            #d = (c**2 - a**2 - b**2)/(2*a**2)
                            #m_coor = pseudo_atom_coor + vec_a * d/a
                            #h_coor2 = h_coor + 2 * (m_coor - h_coor)
                            #h_atoms += [[h_atom_idx, h_coor2]]

                    fea_dict_new['h_atoms'] = h_atoms
                    fea_dict_new['num_neighbor_atom'] = num_neighbor_atom

                elif feature_type == 'Aromatic':
                    v_a = pseudo_atom[1] - pseudo_atom[0]
                    v_b = pseudo_atom[2] - pseudo_atom[1]
                    v_c = np.cross(v_a, v_b)
                    v_n = v_c / np.linalg.norm(v_c)
                    fea_dict_new['v_n'] = v_n
                fea_dict_new['atom_type'] = fea_dict['atom_type']
                if is_protein:
                    fea_dict_new['chain_id'] = fea_dict['chain_id']
                    fea_dict_new['residue_name'] = fea_dict['residue_name']
                    fea_dict_new['residue_num'] = fea_dict['residue_num']
                    fea_dict_new['weight'] = fea_dict['weight']

                PF_coor[feature_type] += [fea_dict_new]

        return PF_coor

    def find_metal_ion(self, m, is_protein=True):
        """
            input:
                m: pybel molecule object
                is_protein: True or False
            output:
                metal_coor: list
        """
        metal_coor = list()
        atoms = m.atoms
        for atom in atoms:
            atomic_num = atom.atomicnum

            if atomic_num not in self.metal_atomic_name_dict:
                continue
            atomic_name = self.metal_atomic_name_dict[atomic_num]
            atom_idx = atom.idx
            atom_idx_group = (atom_idx,)
            coor = np.array(atom.coords)
            if atomic_name not in self.metal_name_dict:
                continue
            pattern_idx = self.metal_name_dict[atomic_name]

            fea_dict = dict()
            fea_dict['atom_type'] = atom.type
            fea_dict['atom_idx_group'] = atom_idx_group
            fea_dict['pattern_idx'] = pattern_idx
            fea_dict['pseudo_atom_coor'] = coor
            if is_protein:
                residue = atom.residue
                chain_id = residue.OBResidue.GetChain()
                fea_dict['chain_id'] = chain_id
                residue_name = residue.name
                fea_dict['residue_name'] = residue_name
                residue_num = residue.OBResidue.GetNum()
                fea_dict['residue_num'] = residue_num
                fea_dict['weight'] = 1.0

            metal_coor += [fea_dict]

        return metal_coor

    def find_mba(self, m, is_protein=False):
        """
            metal binding atom for ligand
            input:
                m: pybel molecule object
                is_protein: True or False
            output:
                mba_coor: list
        """

        mba_coor = list()
        atoms = m.atoms
        for atom in atoms:
            atom_idx = atom.idx
            atomic_num = atom.atomicnum
            if atomic_num != 7 and atomic_num != 8 and atomic_num != 16:
                continue

            # if atom.formalcharge>0:
            #    continue
            # Excluded because protonation state calculation is incomplete

            atom_idx_group = (atom_idx,)
            coor = np.array(atom.coords)
            pattern_idx = self.metal_bind_name_dict[atomic_num]

            fea_dict = dict()
            fea_dict['atom_type'] = atom.type
            fea_dict['atom_idx_group'] = atom_idx_group
            fea_dict['pattern_idx'] = pattern_idx
            fea_dict['pseudo_atom_coor'] = coor
            if is_protein:
                residue = atom.residue
                chain_id = residue.OBResidue.GetChain()
                fea_dict['chain_id'] = chain_id
                residue_name = residue.name
                fea_dict['residue_name'] = residue_name
                residue_num = residue.OBResidue.GetNum()
                fea_dict['residue_num'] = residue_num
                fea_dict['weight'] = 1.0
            mba_coor += [fea_dict]

        return mba_coor

    def draw_ligand(self, m_ligand, PF_dict, output_name, size):
        type_list = self.feature_type_list
        patom_list = list()
        for fea in type_list:
            patoms = list()
            fff = PF_dict[fea]
            for ff in fff:
                atom_idx_group = ff[0]
                for idx in atom_idx_group:
                    patoms += [idx-1]
            patom_list += [patoms]
        m_ligand_noH = Chem.RemoveHs(m_ligand)
        m_ligand_noH.RemoveConformer(0)
        mol_list = [m_ligand_noH]*6
        img = Draw.MolsToGridImage(mol_list, legends=type_list,
                                   highlightAtomLists=patom_list,
                                   subImgSize=size, molsPerRow=2)
        img.save(output_name)

    def box_protein(self, PF_coor_protein, cmin, cmax):
        """
            select PF of PF_coor_protein in the box
            input:
                PF_coor_protein: dict
                cmin : np.array
                cmax : np.array
            output:
                PF_coor_box: dict
        """

        type_list = self.feature_type_list + ['Metal']
        PF_coor_box = dict()
        for fea in type_list:
            if fea not in PF_coor_protein:
                continue
            fea_dict_list = PF_coor_protein[fea]
            PF_coor_box[fea] = list()
            for fea_dict in fea_dict_list:
                pseudo_atom_coor = fea_dict['pseudo_atom_coor']
                if (pseudo_atom_coor > cmax).any():
                    continue
                elif (pseudo_atom_coor < cmin).any():
                    continue
                PF_coor_box[fea] += [fea_dict]
        return PF_coor_box

    def find_template(self, PF_coor_protein, template_ligand_file, include_hydrophobic=False):
        template_dict = dict()
        feature_type_list_receptor = self.feature_type_list + ['Metal']
        for feature_type_receptor in feature_type_list_receptor:
            if not include_hydrophobic and feature_type_receptor == 'Hydrophobic':
                continue
            template_dict[feature_type_receptor] = dict()

        file_format = template_ligand_file.split('.')[-1]
        ms = list(pybel.readfile(file_format, template_ligand_file))
        m_ligand = ms[0]
        ligand_bond_dict = self.get_bond_info(m_ligand)
        PF_dict_ligand = self.find_PF(m_ligand, ligand_bond_dict)

        num_model = len(ms)
        if num_model > 1:
            num_model = num_model - 1
        PF_coor_ligand_dict = dict()
        interaction_dict = dict()
        model_idx = 0
        m_ligand = ms[model_idx]
        PF_coor_ligand = self.find_PF_coor(m_ligand, PF_dict_ligand, ligand_bond_dict)
        mba_coor = self.find_mba(m_ligand, is_protein=False)
        PF_coor_ligand['MBA'] = mba_coor

        PF_coor_ligand_dict[model_idx] = PF_coor_ligand
        interaction = self.find_interaction(PF_coor_protein, PF_coor_ligand,
                                            ligand_bond_dict)
        interaction_dict[model_idx] = interaction
        total_rec_pf_dict = self.count_interaction_receptor(interaction_dict)
        rec_pf_dict = total_rec_pf_dict[0]
        for feature_type_receptor in rec_pf_dict.keys():
            if not include_hydrophobic and feature_type_receptor == 'Hydrophobic':
                continue
            rec_pf = rec_pf_dict[feature_type_receptor]
            atom_idx_group_list = rec_pf.keys()
            for atom_idx_group in atom_idx_group_list:
                if atom_idx_group not in template_dict[feature_type_receptor]:
                    template_dict[feature_type_receptor][atom_idx_group] = 0
                template_dict[feature_type_receptor][atom_idx_group] += 1
        return template_dict

    def select_template_protein(self, PF_coor_protein, template_dict,
                                cut_num=0):
        """
            select PF of PF_coor_protein in template
            input:
                PF_coor_protein: dict
                template_dict: dict
            output:
                PF_coor_info: dict
        """
        type_list = self.feature_type_list + ['Metal']
        PF_coor_info = dict()
        for fea in type_list:
            if fea not in PF_coor_protein:
                continue
            if fea not in template_dict:
                continue
            fea_dict_list = PF_coor_protein[fea]
            PF_coor_info[fea] = list()
            for fea_dict in fea_dict_list:
                atom_idx_group = fea_dict['atom_idx_group']
                if atom_idx_group not in template_dict[fea]:
                    continue
                num_in_template = template_dict[fea][atom_idx_group]
                if num_in_template < cut_num:
                    continue
                PF_coor_info[fea] += [fea_dict]
        return PF_coor_info

    def find_interaction(self, PF_coor_protein, PF_coor_ligand,
                         ligand_bond_dict):
        """
            find pharmacophoric interaction
            input:
                PF_coor_protein: dict
                PF_coor_ligand: dict
                ligand_bond_dict: dict
            output:
                interaction: dict
        """
        type_list = self.feature_type_list + ['MBA']
#        type_list = self.feature_type_list

        coor_ligand_list = list()
        for fea_ligand in type_list:
            if fea_ligand not in PF_coor_ligand:
                continue
            fea_dict_list_ligand = PF_coor_ligand[fea_ligand]
            for fea_dict_ligand in fea_dict_list_ligand:
                coor_ligand = fea_dict_ligand['pseudo_atom_coor']
                coor_ligand_list.append(coor_ligand)
        coor_ligand_list = np.array(coor_ligand_list)
        cmin = coor_ligand_list.min(axis=0) - self.cutoff
        cmax = coor_ligand_list.max(axis=0) + self.cutoff

        PF_coor_protein_box = self.box_protein(PF_coor_protein, cmin, cmax)

        interaction = dict()
        for fea_ligand in type_list:
            if fea_ligand not in interaction:
                interaction[fea_ligand] = dict()
            if fea_ligand not in PF_coor_ligand:
                continue
            fea_dict_list_ligand = PF_coor_ligand[fea_ligand]
            cfea_list = self.complementary_feature[fea_ligand]
            for fea_dict_ligand in fea_dict_list_ligand:
                idx_ligand = fea_dict_ligand['atom_idx_group']
                pattern_ligand = fea_dict_ligand['pattern_idx']
                coor_ligand = fea_dict_ligand['pseudo_atom_coor']
                for cfea_t in cfea_list:
                    fea_protein = cfea_t[0]
                    if fea_protein not in interaction[fea_ligand]:
                        interaction[fea_ligand][fea_protein] = list()
                    pair_cutoff = cfea_t[1]
                    if fea_protein not in PF_coor_protein_box:
                        continue
                    fea_dict_list_protein = PF_coor_protein_box[fea_protein]
                    for fea_dict_protein in fea_dict_list_protein:
                        idx_protein = fea_dict_protein['atom_idx_group']
                        pattern_protein = fea_dict_protein['pattern_idx']
                        coor_protein = fea_dict_protein['pseudo_atom_coor']

                        r_lp = coor_protein - coor_ligand
                        dist_lp = np.linalg.norm(r_lp)
#                        if dist_lp > self.cutoff:
                        if dist_lp > pair_cutoff:
                            continue

                        interact_ij = dict()
                        interact_ij['idx_ligand'] = idx_ligand
                        interact_ij['idx_protein'] = idx_protein
                        interact_ij['pattern_ligand'] = pattern_ligand
                        interact_ij['pattern_protein'] = pattern_protein

                        interact_ij['dist_lp'] = dist_lp
                        interact_ij['chain_id'] = fea_dict_protein['chain_id']
                        interact_ij['residue_name'] = fea_dict_protein['residue_name']
                        interact_ij['residue_num'] = fea_dict_protein['residue_num']
                        interact_ij['weight'] = fea_dict_protein['weight']

                        if fea_ligand == 'HBD' and fea_protein == 'HBA':
                            num_neighbor_atom = fea_dict_ligand['num_neighbor_atom']
                            donor_hydrogen_list = fea_dict_ligand['h_atoms']
                            if num_neighbor_atom <= 1:
                                theta = 0
                                interact_ij['theta'] = theta
                            else:
                                dist_ah = 100.0
                                for donor_hydrogen in donor_hydrogen_list:
                                    hydrogen_coor0 = donor_hydrogen[1]
                                    dist_ah0 = np.linalg.norm(
                                        hydrogen_coor0 - coor_protein)
                                    if dist_ah0 < dist_ah:
                                        dist_ah = dist_ah0
                                        hydrogen_coor = hydrogen_coor0
                                dist_dh = np.linalg.norm(
                                    hydrogen_coor - coor_ligand)
                                theta = cal_angle_from_points(
                                    dist_lp, dist_dh, dist_ah)
                                if theta < 0.6*np.pi:
                                    continue
                                interact_ij['theta'] = theta

                        elif fea_ligand == 'HBA' and fea_protein == 'HBD':
                            num_neighbor_atom = fea_dict_protein['num_neighbor_atom']
                            donor_hydrogen_list = fea_dict_protein['h_atoms']
                            if num_neighbor_atom <= 1:
                                theta = 0
                                interact_ij['theta'] = theta
                            else:
                                dist_ah = 100.0
                                for donor_hydrogen in donor_hydrogen_list:
                                    hydrogen_coor0 = donor_hydrogen[1]
                                    dist_ah0 = np.linalg.norm(
                                        hydrogen_coor0 - coor_ligand)
                                    if dist_ah0 < dist_ah:
                                        dist_ah = dist_ah0
                                        hydrogen_coor = hydrogen_coor0
                                dist_dh = np.linalg.norm(
                                    hydrogen_coor - coor_protein)
                                theta = cal_angle_from_points(
                                    dist_lp, dist_dh, dist_ah)
                                if theta < 0.6*np.pi:
                                    continue
                                interact_ij['theta'] = theta

                        elif fea_ligand == 'Aromatic' and fea_protein == 'Cation':
                            n_vector_ligand = fea_dict_ligand['v_n']
                            n_vector_lp = r_lp/dist_lp
                            theta = np.abs(cal_angle_from_vectors(
                                           n_vector_ligand, n_vector_lp))
                            if theta > np.pi/2:
                                theta = np.pi-theta
                            if theta > 0.5*np.pi/2:
                                continue
                            interact_ij['theta'] = theta

                        elif fea_ligand == 'Cation' and fea_protein == 'Aromatic':
                            n_vector_protein = fea_dict_protein['v_n']
                            n_vector_lp = r_lp/dist_lp
                            theta = np.abs(cal_angle_from_vectors(
                                           n_vector_protein, n_vector_lp))
                            if theta > np.pi/2:
                                theta = np.pi-theta
                            if theta > 0.5*np.pi/2:
                                continue
                            interact_ij['theta'] = theta

                        elif fea_ligand == 'Aromatic' and fea_protein == 'Aromatic':
                            n_vector_ligand = fea_dict_ligand['v_n']
                            n_vector_protein = fea_dict_protein['v_n']
                            n_vector_lp = r_lp/dist_lp
                            theta = cal_angle_from_vectors(
                                n_vector_ligand, n_vector_protein)
                            if theta > np.pi/2:
                                theta = np.pi-theta
                            interact_ij['theta'] = theta

                            alpha = np.abs(cal_angle_from_vectors(
                                           n_vector_ligand, n_vector_lp))
#                            dist_v = dist_lp * np.cos(alpha)
#                            dist_h = dist_lp * np.sin(alpha)
#                            print(np.sqrt(dist_v*dist_v+dist_h*dist_h), dist_lp)
#                            interact_ij['dist_v'] = dist_v
#                            interact_ij['dist_h'] = dist_h
                            if theta > 0.4*np.pi/2 and theta < 0.8*np.pi/2:
                                continue
                            interact_ij['alpha'] = alpha

                        elif fea_ligand == 'MBA' and fea_protein == 'Metal':
                            if pattern_ligand not in self.metal_interaction_cutoff:
                                continue
                            if pattern_protein not in self.metal_interaction_cutoff[pattern_ligand]:
                                continue
                            metal_cutoff = self.metal_interaction_cutoff[
                                pattern_ligand][pattern_protein]
                            if dist_lp > metal_cutoff:
                                continue
                        interaction[fea_ligand][fea_protein] += [interact_ij]

        if 'MBA' in interaction:
            if 'Metal' in interaction['MBA']:
                metal_interaction = interaction['MBA']['Metal']
                d_list = list()
                d_list = np.array([x['dist_lp'] for x in metal_interaction])
                idx_sorted = np.argsort(d_list)

                metal_interaction_new = list()
                batom_idx = list()
                for idx in idx_sorted:
                    interact_ij = metal_interaction[idx]
                    idx_ligand = interact_ij['idx_ligand'][0]
#                    dist_lp = interact_ij['dist_lp']
                    neighbor_atoms_idx = ligand_bond_dict[idx_ligand]
                    neighbor_check = False
                    for neighbor_atom_idx in neighbor_atoms_idx:

                        if neighbor_atom_idx in batom_idx:
                            neighbor_check = True
                    if not neighbor_check:
                        metal_interaction_new += [interact_ij]
                        batom_idx += [idx_ligand]

                interaction['MBA']['Metal'] = metal_interaction_new

        return interaction

    def write_PF(self, PF_coor_dict, out_file, is_protein=False):
        """
            write pharmacophoric feature
            input:
                PF_coor_dict: dict
                out_file: filename, str
                is_protein: bool
        """

        if is_protein:
            feature_type_list = self.feature_type_list + ['Metal']
        else:
            feature_type_list = self.feature_type_list + ['MBA']
        fp = open(out_file, 'w')
        line_out = 'feature_type:atom_idx_group:pattern_idx:pseudo_atom_coor'
        line_out += ':atom_type:etc'
        if is_protein:
            line_out += ':chain_id:residue_name:residue_num:weight'
        line_out += '\n'
        fp.write(line_out)
        model_idx_list = PF_coor_dict.keys()
        num_model = len(model_idx_list)
        for model_idx in model_idx_list:
            PF_coor = PF_coor_dict[model_idx]
            if num_model > 1:
                line_out = 'MODEL %d\n' % (model_idx + 1)
                fp.write(line_out)
            for feature_type in feature_type_list:
                if feature_type not in PF_coor:
                    continue
                fea_dict_list = PF_coor[feature_type]
                for fea_dict in fea_dict_list:
                    line_out = '%s' % feature_type
#                    line_out += ':%s' % (str(fea_dict['atom_idx_group']))
                    idx_line = ''
                    for atom_idx in fea_dict['atom_idx_group']:
                        idx_line += '%d,' % atom_idx
                    line_out += ':(%s)' % idx_line.strip(',')

                    line_out += ':%s' % (fea_dict['pattern_idx'])
                    coor = fea_dict['pseudo_atom_coor']
                    line_out += ':(%.3f,%.3f,%.3f)' % (
                        coor[0], coor[1], coor[2])
                    line_out += ':%s' % (fea_dict['atom_type'])

                    if 'h_atoms' in fea_dict:
                        hatoms = fea_dict['h_atoms']
                        hatom_line = ''
                        if 'num_neighbor_atom' in fea_dict:
                            num_neighbor_atom = fea_dict['num_neighbor_atom']
                            hatom_line += '%d;' %num_neighbor_atom
                        for hatom in hatoms:
                            line_f = '%d=(%.3f,%.3f,%.3f);'
                            hatom_line += line_f % (hatom[0], hatom[1][0],
                                                    hatom[1][1], hatom[1][2])
                        line_out += ':%s' % hatom_line.strip(';')

                    elif 'v_n' in fea_dict:
                        coor_v_n = fea_dict['v_n']
                        line_out += ':(%.3f,%.3f,%.3f)' % (
                            coor_v_n[0], coor_v_n[1], coor_v_n[2])

                    else:
                        line_out += ':'

                    if is_protein:
                        line_out += ':%s' % (fea_dict['chain_id'])
                        line_out += ':%s' % (fea_dict['residue_name'])
                        line_out += ':%d' % (fea_dict['residue_num'])
                        line_out += ':%.3f' % (fea_dict['weight'])

                    line_out += '\n'
                    fp.write(line_out)
        fp.close()

    def read_PF(self, feature_file, is_protein=False):
        PF_coor_dict = dict()
        fp = open(feature_file)
        lines = fp.readlines()
        fp.close()

        model_idx = 0
        PF_coor_dict[model_idx] = dict()

        for line in lines:
            if line.startswith('feature_type'):
                # title = line.strip().split(':')
                continue

            if line.startswith('MODEL'):
                model_idx = int(line.strip().split()[1]) - 1
                fea_dict = dict()
                if model_idx not in PF_coor_dict:
                    PF_coor_dict[model_idx] = dict()
                continue
            lis = line.strip().split(':')
            if len(lis) < 1:
                continue
            fea_dict = dict()
            feature_type = lis[0]
            atom_idx_group = lis[1]
            pattern_idx = lis[2]
            pseudo_atom_coor = lis[3]
            atom_type = lis[4]
            etc = lis[5]
            if feature_type not in PF_coor_dict[model_idx]:
                PF_coor_dict[model_idx][feature_type] = list()

            fea_dict['feature_type'] = feature_type
            aline = atom_idx_group.lstrip('(').rstrip(')')
            fea_dict['atom_idx_group'] = np.array(aline.split(','), dtype=int)
            fea_dict['pattern_idx'] = int(pattern_idx)
            coor = pseudo_atom_coor.lstrip('(').rstrip(')').split(',')
            fea_dict['pseudo_atom_coor'] = np.array(coor, dtype=np.float32)
            fea_dict['atom_type'] = atom_type
            if feature_type == 'HBD':
                hline_list = etc.split(';')
                h_list = list()
                num_neighbor_atom = float(hline_list[0])
                fea_dict['num_neighbor_atom'] = num_neighbor_atom
                for hline in hline_list[1:]:
                    aa = hline.split('=')
                    h_idx = int(aa[0])
                    h_coor = aa[1].lstrip('(').rstrip(')').split(',')
                    h_list += [[h_idx, np.array(h_coor, dtype=np.float32)]]
                    fea_dict['h_atoms'] = h_list
            if feature_type == 'Aromatic':
                v_n = etc.lstrip('(').rstrip(')').split(',')
                fea_dict['v_n'] = np.array(v_n, dtype=np.float32)
            if is_protein:
                chain_id = lis[6]
                residue_name = lis[7]
                residue_num = lis[8]
                fea_dict['chain_id'] = chain_id
                fea_dict['residue_name'] = residue_name
                fea_dict['residue_num'] = int(residue_num)
                if len(lis) >= 10:
                    weight = lis[9]
                    fea_dict['weight'] = float(weight)
                else:
                    fea_dict['weight'] = 1.0

            PF_coor_dict[model_idx][feature_type] += [fea_dict]

        return PF_coor_dict

    def write_interaction(self, interaction_dict, out_file):
        """
            write pharmacophoric interaction
            input:
                interaction_dict: dict()
                out_file: str
        """
        fp = open(out_file, 'w')
        line_out = 'feature_type_ligand:feature_type_protein'
        line_out += ':atom_idx_group_ligand:atom_idx_group_protein'
        line_out += ':pattern_ligand:pattern_protein'
        line_out += ':chain_id:residue_name:residue_num'
        line_out += ':dist_lp:theta:alpha:weight\n'
        fp.write(line_out)

        model_idx_list = interaction_dict.keys()
        num_model = len(model_idx_list)
        for model_idx in model_idx_list:
            interaction = interaction_dict[model_idx]
            if num_model > 1:
                line_out = 'MODEL %d\n' % (model_idx + 1)
                fp.write(line_out)

            for feature_type_ligand in interaction.keys():
                inter_i = interaction[feature_type_ligand]
                for feature_type_protein in inter_i.keys():
                    inter_ij_list = inter_i[feature_type_protein]
                    for inter_ij in inter_ij_list:
                        line_out = '%s' % feature_type_ligand
                        line_out += ':%s' % feature_type_protein

                        idx_line = ''
                        for idx in inter_ij['idx_ligand']:
                            idx_line += '%d,' % (idx)
                        line_out += ':(%s)' % (idx_line.strip(','))

                        idx_line = ''
                        for idx in inter_ij['idx_protein']:
                            idx_line += '%d,' % (idx)
                        line_out += ':(%s)' % (idx_line.strip(','))
                        line_out += ':%s' % (inter_ij['pattern_ligand'])
                        line_out += ':%s' % (inter_ij['pattern_protein'])

                        line_out += ':%s' % (inter_ij['chain_id'])
                        line_out += ':%s' % (inter_ij['residue_name'])
                        line_out += ':%d' % (inter_ij['residue_num'])
                        line_out += ':%.3f' % (inter_ij['dist_lp'])
                        if 'theta' in inter_ij:
                            line_out += ':%.3f' % (inter_ij['theta'])
                        else:
                            line_out += ':'
                        if 'alpha' in inter_ij:
                            line_out += ':%.3f' % (inter_ij['alpha'])
                        else:
                            line_out += ':'
                        line_out += ':%.3f' % (inter_ij['weight'])

                        line_out += '\n'
                        fp.write(line_out)

        fp.close()

    def count_interaction_receptor(self, interaction_dict):

        total_rec_pf_dict = dict()
        model_idx_list = interaction_dict.keys()
        feature_type_list_receptor = self.feature_type_list + ['Metal']
        feature_type_list_ligand = self.feature_type_list + ['MBA']

        for model_idx in model_idx_list:
            line_out_title = 'Model_idx'
            interaction = interaction_dict[model_idx]

            rec_pf_dict = dict()
            for feature_type_ligand in feature_type_list_ligand:
                line_out_title += ' %s' % (feature_type_ligand)
                if feature_type_ligand not in interaction:
                    continue
                inter_i = interaction[feature_type_ligand]
                for feature_type_protein in feature_type_list_receptor:
                    if feature_type_protein not in rec_pf_dict:
                        rec_pf_dict[feature_type_protein] = dict()
                    if feature_type_protein not in inter_i.keys():
                        continue
                    inter_ij_list = inter_i[feature_type_protein]
                    for inter_ij in inter_ij_list:
                        idx_protein = inter_ij['idx_protein']
                        if idx_protein not in rec_pf_dict[feature_type_protein]:
                            rec_pf_dict[feature_type_protein][idx_protein] = 0
                        rec_pf_dict[feature_type_protein][idx_protein] += 1
            total_rec_pf_dict[model_idx] = rec_pf_dict

        return total_rec_pf_dict

    def print_interaction_receptor(self, total_rec_pf_dict):
        line_out_total = ''
        model_idx_list = total_rec_pf_dict.keys()
        feature_type_list_receptor = self.feature_type_list + ['Metal']

        line_out_title = 'Model_idx'
        for feature_type_receptor in feature_type_list_receptor:
            line_out_title += ' %s' % (feature_type_receptor)
        line_out_total += line_out_title + '\n'

        for model_idx in model_idx_list:
            line_out_value = '%d' % (model_idx + 1)
            rec_pf_dict = total_rec_pf_dict[model_idx]
            for feature_type_protein in feature_type_list_receptor:
                if feature_type_protein in rec_pf_dict:
                    count = len(rec_pf_dict[feature_type_protein].keys())
                else:
                    count = 0
                line_out_value += ' %d' % (count)
            line_out_total += line_out_value + '\n'

        return line_out_total

    def count_interaction_ligand(self, interaction_dict):
        total_lig_pf_dict = dict()
        model_idx_list = interaction_dict.keys()
        feature_type_list = self.feature_type_list + ['MBA']

        for model_idx in model_idx_list:
            interaction = interaction_dict[model_idx]
            lig_pf_dict = dict()
            for feature_type_ligand in feature_type_list:
                if feature_type_ligand not in interaction:
                    continue
                inter_i = interaction[feature_type_ligand]
                lig_pf_dict[feature_type_ligand] = dict()
                for feature_type_protein in inter_i.keys():
                    inter_ij_list = inter_i[feature_type_protein]
                    for inter_ij in inter_ij_list:
                        idx_ligand = inter_ij['idx_ligand']
                        if idx_ligand not in lig_pf_dict[feature_type_ligand]:
                            lig_pf_dict[feature_type_ligand][idx_ligand] = 0
                        lig_pf_dict[feature_type_ligand][idx_ligand] += 1
            total_lig_pf_dict[model_idx] = lig_pf_dict

        return total_lig_pf_dict

    def print_interaction_ligand(self, total_lig_pf_dict):
        line_out_total = ''
        model_idx_list = total_lig_pf_dict.keys()
        feature_type_list = self.feature_type_list + ['MBA']

        line_out_title = 'Model_idx'
        for feature_type_ligand in feature_type_list:
            line_out_title += ' %s' % (feature_type_ligand)
        line_out_total += line_out_title + '\n'

        for model_idx in model_idx_list:
            line_out_value = '%d' % (model_idx + 1)
            lig_pf_dict = total_lig_pf_dict[model_idx]
            for feature_type_ligand in feature_type_list:
                if feature_type_ligand in lig_pf_dict:
                    count = len(lig_pf_dict[feature_type_ligand].keys())
                else:
                    count = 0
                line_out_value += ' %d' % (count)
            line_out_total += line_out_value + '\n'

        return line_out_total

    def count_interaction(self, interaction_dict, use_weight_pis=True,
                          include_hydrophobic=True):
        """
            input:
                interaction_dict: dict()
            output:
                total_count_dict: dict()
        """

        total_count_dict = dict()

        feature_type_list = self.feature_type_list + ['MBA']
        model_idx_list = interaction_dict.keys()
        for model_idx in model_idx_list:
            interaction = interaction_dict[model_idx]
            count_dict = dict()
            w_count_dict = dict()
            for feature_type_ligand in feature_type_list:
                cfea_list = self.complementary_feature[feature_type_ligand]
                for cfea in cfea_list:
                    feature_type_protein = cfea[0]

                    key = '%s:%s' % (feature_type_ligand, feature_type_protein)
                    if key not in count_dict:
                        count_dict[key] = 0
                        w_count_dict[key] = 0
                    if feature_type_ligand not in interaction:
                        continue
                    inter_i = interaction[feature_type_ligand]
                    if feature_type_protein not in inter_i:
                        continue
                    if ((not include_hydrophobic)
                            and feature_type_protein == 'Hydrophobic'):
                        continue

                    inter_ij_list = inter_i[feature_type_protein]
                    for inter_ij in inter_ij_list:
                        count_dict[key] += 1
                        weight = inter_ij['weight']
                        w_count_dict[key] += weight

            score = 0
            for key in self.interaction_feature:
                if use_weight_pis:
                    score += w_count_dict[key]*self.weight_pis[key]
                else:
                    score += w_count_dict[key]
            if use_weight_pis:
                score += self.bias_pis

            count_dict['Score'] = score

            total_count_dict[model_idx] = count_dict

        return total_count_dict

    def print_interaction(self, total_count_dict, total_count_info_dict,
                          use_pinfo, print_option=1):
        """
            input:
                total_count_dict: from count_interaction
                print_option:
                    if 1, print interction in one line for one model.
                    if 2, print interaction in multi lines.
            output:
                line_out_total: str
        """

        line_out_total = ''

        model_idx_list = total_count_dict.keys()
        num_model = len(model_idx_list)
        if print_option == 1:
            line_out_title = 'Model_idx PIscore'
            if use_pinfo:
                line_out_title += ' Pinfo'
            for key in self.interaction_feature:
                line_out_title += ' %s' % (key)
            line_out_total += line_out_title + '\n'

        for model_idx in model_idx_list:
            count_dict = total_count_dict[model_idx]
            if num_model > 1 and print_option == 2:
                line_out = 'MODEL %d' % (model_idx + 1)
                line_out_total += line_out + '\n'
            if print_option == 1:
                line_out_value = '%d' % (model_idx + 1)
                line_out_value += ' %.4f' % (count_dict['Score'])
                if use_pinfo:
                    count_info_dict = total_count_info_dict[model_idx]
                    line_out_value += ' %.4f' % (count_info_dict['Score'])

            for key in self.interaction_feature:
                count = count_dict[key]
                if print_option == 1:
                    line_out_value += ' %d' % (count)

                if count > 0 and print_option == 2:
                    line_out = '%s %d\n' % (key, count)
                    line_out_total += line_out

            if print_option == 2:
                line_out = 'PIscore: %.4f\n' % (count_dict['Score'])
                if use_pinfo:
                    count_info_dict = total_count_info_dict[model_idx]
                    line_out += 'Pinfo %.4f\n' % (count_info_dict['Score'])

                line_out_total += line_out

            if print_option == 1:
                line_out_total += line_out_value + '\n'

        return line_out_total

    def read_dock_config(self, dock_config):
        fp = open(dock_config)
        lines = fp.readlines()
        fp.close()

        for line in lines:
            lis = line.strip().split('=')
            if lis[0] == 'center_x':
                center_x = float(lis[1])
            elif lis[0] == 'center_y':
                center_y = float(lis[1])
            elif lis[0] == 'center_z':
                center_z = float(lis[1])
            elif lis[0] == 'size_x':
                size_x = float(lis[1])
            elif lis[0] == 'size_y':
                size_y = float(lis[1])
            elif lis[0] == 'size_z':
                size_z = float(lis[1])

        center = np.array((center_x, center_y, center_z), dtype=float)
        size = np.array((size_x, size_y, size_z), dtype=float)
        return center, size

    def cal_piscore(self, ligand_file, pf_ligand_file, interaction_file):

        ms = list(pybel.readfile('pdb', ligand_file))
        num_model = len(ms)
        if num_model > 1:
            num_model = num_model - 1
        if num_model < 1:
            return None

        m_ligand = ms[0]

        PF_coor_ligand_dict = dict()
        interaction_dict = dict()
        if self.use_pinfo:
            interaction_info_dict = dict()
        ligand_bond_dict = self.get_bond_info(m_ligand)
        PF_dict_ligand = self.find_PF(m_ligand, ligand_bond_dict)
        for model_idx in range(num_model):
            m_ligand = ms[model_idx]
            PF_coor_ligand = self.find_PF_coor(m_ligand, PF_dict_ligand, ligand_bond_dict)
            mba_coor = self.find_mba(m_ligand, is_protein=False)
            PF_coor_ligand['MBA'] = mba_coor
            PF_coor_ligand_dict[model_idx] = PF_coor_ligand
            interaction = self.find_interaction(self.PF_coor_protein,
                                                PF_coor_ligand,
                                                ligand_bond_dict)

            interaction_dict[model_idx] = interaction
            if self.use_pinfo:
                interaction_info = self.find_interaction(self.PF_coor_info,
                                                         PF_coor_ligand,
                                                         ligand_bond_dict)
                interaction_info_dict[model_idx] = interaction_info

        if pf_ligand_file is not None:
            self.write_PF(PF_coor_ligand_dict, pf_ligand_file)

        if interaction_file is not None:
            self.write_interaction(interaction_dict, interaction_file)

        total_count_dict = self.count_interaction(interaction_dict,
                                                  use_weight_pis=True,
                                                  include_hydrophobic=True)
        if self.use_pinfo:
            total_count_info_dict = self.count_interaction(
                interaction_info_dict,
                use_weight_pis=False,
                include_hydrophobic=True)
        else:
            total_count_info_dict = None
        return total_count_dict, total_count_info_dict

def set_pifinder(args):

    piscore_receptor = args.piscore_receptor
    pf_receptor = args.pf_receptor
    pinfo_ligand = args.pinfo_ligand
    pf_receptor_info = args.pf_receptor_info
    pi_cutoff = args.pi_cutoff
    include_hydrophobic = args.include_hydrophobic
    dock_config = args.dock_config

    params = dict()
    weight_pis = {
        'HBD:HBA': 0.129,
        'HBA:HBD': 0.129,
        'Anion:Cation': 0.31,
        'Cation:Anion': 0.31,
        'Cation:Aromatic': 0.039,
        'Hydrophobic:Hydrophobic': 0.00585,
        'Aromatic:Cation': 0.039,
        'Aromatic:Aromatic': 0.064,
        'MBA:Metal': 1.0
    }
    params['weight_pis'] = weight_pis
    params['bias_pis'] = 4.0
    params['cutoff'] = pi_cutoff
    if dock_config is not None:
        params['dock_config'] = dock_config
    params['include_hydrophobic'] = include_hydrophobic
    params['piscore_receptor'] = piscore_receptor
    params['pf_receptor'] = pf_receptor
    params['pinfo_ligand'] = pinfo_ligand
    params['pf_receptor_info'] = pf_receptor_info

    pharma = Pharmacophore(params)

    return pharma

def main():

    import sys
    import argparse
    title_line = 'Fixer for ligand pdb which is converted from pdbqt'
    parser = argparse.ArgumentParser(description=title_line)
    parser.add_argument('-r', '--piscore_receptor', required=False,
                        default=None,
                        help='input receptor pdb file')
    parser.add_argument('-p', '--pf_receptor', required=False,
                        default=None,
                        help='output for pharmacophoric feature of receptor')
    parser.add_argument('-t', '--pinfo_ligand', required=False,
                        default=None, help='pinfo_ligand')
    parser.add_argument('-u', '--pf_receptor_info', required=False,
                        default=None,
                        help='output for template feature of receptor')
    parser.add_argument('-v', '--dock_config', type=str, required=False,
                        default=None, help='dock_config_file')
    parser.add_argument('-m', '--pi_cutoff', type=float, required=False,
                        default=6.5,
                        help='pharmacophoric interaction cutoff distance')
    parser.add_argument('-l', '--ligand_file', required=True,
                        help='input ligand pdb file')
    parser.add_argument('-q', '--pf_ligand_file', required=False,
                        default=None,
                        help='write pharmacophoric feature of ligand')
    parser.add_argument('--include_hydrophobic', action='store_true',
                        required=False,
                        help='include hydrophobic feature for template')
    parser.add_argument('-i', '--interaction_file', required=False,
                        default=None, help='write interaction')
    parser.add_argument('--print_option', type=int, required=False, default=0,
                        help='print_option 0:None, 1: one_line, 2: multi_line')

#    draw_ligand = False
    args = parser.parse_args()
    if args.piscore_receptor is None and args.pf_receptor is None:
        parser.print_usage()
        print('error piscore_receptor and pf_receptor are None')
        sys.exit()

    pharma = set_pifinder(args)
    use_pinfo = pharma.use_pinfo

    ligand_file = args.ligand_file
    pf_ligand_file = args.pf_ligand_file
    interaction_file = args.interaction_file
    print_option = args.print_option

    result = pharma.cal_piscore(ligand_file, pf_ligand_file, interaction_file)
    total_count_dict, total_count_info_dict = result

#    if draw_ligand:
#        size = (200, 200)
#        fig_dir = 'fig'
#        output_name = fig_dir + '/%s.png' % (ligand_file[:-3])
#        m_rdkit_ligand = Chem.MolFromPDBFile(ligand_file, removeHs=True)
#        pharma.draw_ligand(m_rdkit_ligand, PF_dict_ligand, output_name, size)

    if print_option > 0:
        lines_count = pharma.print_interaction(total_count_dict,
                                               total_count_info_dict,
                                               use_pinfo,
                                               print_option=print_option)
        print(lines_count.rstrip())

if __name__ == '__main__':
    main()


![Callysto.ca Banner](https://github.com/callysto/curriculum-notebooks/blob/master/callysto-notebook-banner-top.jpg?raw=true)

from IPython.display import HTML
hide_me = ''
HTML('''
<form action="javascript:code_toggle()"><input style="opacity:1" type="submit" value="Click here to toggle on/off the raw code."></form>''')

hide_me

from ipywidgets import interact
import ipywidgets as widgets
import IPython

import matplotlib.pyplot as plt
import numpy as np
import math

import plotly as py
import plotly.graph_objs as go

import pylab

from IPython.display import Image, HTML, YouTubeVideo

# Data Representation in Graphs

### Grade 8 curriculum

Data plays an ever-increasing role in our lives. Like it or not, we are faced with numerical information every day, and we use it to make decisions. Should I be glad that 9 out of 10 dentists recommend my toothpaste? What about the 10th? A new study says that going for a run at 5 a.m. every morning reduces my risk of catching some terrible disease by 15%. Is it worth getting out of bed?

It's often hard to find meaning in data if it's just a bunch of numbers on a page, so we make that easier by using graphs. Graphs take data and turn them into pictures—bars, lines, circles, and more. But not all graphs are created equal; some do their jobs better than others. A good graph is a perfect tool for understanding a problem. A bad graph can be confusing, or in some cases, intentionally misleading.

Graphs are used every day by news media, politicians, and scientists to convey information. Some use them well; some do not. In this notebook, we'll explore good and bad examples of graphs. By working through the examples and exercises in this notebook, you'll learn:

- how to decide which type of graph is best for a given set of data;
- how to identify flawed or misleading graphs;
- how some of those flaws can be corrected; and
- most importantly, how to read a graph and apply its meaning to your everyday life.

*Many of the examples of bad graphs you'll find below are from the media (and one source in particular). This notebook isn't trying to criticize these sources. They just happen to have given us a lot of examples to choose from.*

## What makes a good graph?

First and most importantly, a graph should allow a reader, at a glance, to understand the information it's meant to convey. A good graph is like a good movie scene; if it's set up well, you can tell exactly what you're supposed to know. Some basic parts of a successful graph are:

1. A title
2. Proper labels
3. Axes that start at zero (if numerical)
4. Percentages that add to 100%
5. Easy to read
6. Use of colours, *as long as they are meaningful* and not just for show

*By the way: **axes** (ACK-sees) are the reference lines on a graph. They're where you find the names of the categories (usually at the bottom) and the number scale (usually on the left).  One of these lines is called an **axis** (ACK-sis).*

For a quick overview of different types of graphs and simple examples, you might find this
[Math Is Fun](https://www.mathsisfun.com/data/pictographs.html) article useful. We'll look at some of these kinds of graphs below. You'll notice many of them are eye-catching, and they also convey information really well.

One of the places you'll find a lot of graphs is in political coverage. The media (and many of their readers/viewers) love a good "horse race". For example, this [CBC federal poll tracking article](http://www.cbc.ca/news/politics/poll-tracker-federal-poll-averages-and-seat-projections-1.4171977) uses almost every type of graph you'll find in this notebook.

We'll also explore how a graph can be used to [mislead someone](http://teachersinstitute.yale.edu/curriculum/units/2008/6/08.06.06.x.html). We hope this notebook will help you learn how to avoid using misleading graphs, as well as how to avoid being misled yourself.

There's even a [wall of shame](http://bcuchta.com/wall_of_shame/) with some of the worst graphs and charts!

## Let's look at bar graphs

### What is a bar graph?

A bar graph is a graph where data is separated into categories, and those categories are shown as bars with different heights. It's a very useful graph, but it can also easily be misleading.

![picture](./images/bar-graph-fruit.svg)

from [Math is Fun](https://mathsisfun.com/data/images/bar-graph-fruit.svg)

### When are bar graphs good to use?

Bar graphs can be used in many ways, but they usually show one piece of information collected from many groups or categories. For example, they might show the number of hours worked by people in different age groups, or how many grey shirts each girl in a class owns.

### What are some ways to misuse bar graphs?

1. **Make the scale on the graph start above zero.** This makes small differences between bars look much bigger than they really are.
2. **Change the width of the bars to make one category look more important.** This gives one bar more area, which looks like more data.
3. **Remove the space between the bars** (that's a **histogram**). Histograms are used for a different kind of data set, and so they are read in a different way.

Here's an example of a poorly made bar graph.  It shows the total welfare (support money) received by people in the US from 2009 to 2011.  Each year is divided into 4 three-month pieces called **quarters**.

![picture](./images/fnc-an-20120809-welfarechart.jpg)

from [MediaMatters](https://www.mediamatters.org/fox-news/today-dishonest-fox-charts-government-aid-edition)

What makes this a bad bar graph?
1. Their scale starts at 94 million insead of 0.
2. The bars are in 3D, making their values harder to read.
3. Their y-axis has 8 labels, but there are 10 levels on the graph (including the top and bottom).

Whoever made this graph probably wanted the viewer to think welfare in the US is rising faster than it really is.  Now, let's ask ourselves:

- What can we change to make this a good graph?
- How does it look after we make these changes?
- Why might the original creators not want to use our new graph?

One way we can improve this graph is by changing where its scale starts.  Play with the slider below to see how the graph looks with different scales.

*Slide the slider below to change the starting point of the $y$-axis. The initial position corresponds to the graph in the image above. As you move the slider to the left, the starting point for the $y$-axis is reduced to zero.*

*Warning: This graph is slow to respond, please be patient with it.*

hide_me

columns = ['09-Q1', '09-Q2', '09-Q3', '09-Q4', '10-Q1', '10-Q2','10-Q3', '10-Q4', '11-Q1', '11-Q2']
#fig, ax = plt.subplots()
def plot(yaxis=94):
    y = [97, 98, 99, 101, 104, 105, 106, 107, 107.5, 108]
    x = np.arange(len(y))
    fig, ax = plt.subplots(figsize=(10,4))
    ax.bar(x, y, width=0.5)
    ax.set_xticks(x)
    ax.set_xticklabels(columns)
    ax.set_ylim((yaxis,110))
    ax.set_title("Federal Welfare Received in the US")

interact(plot, yaxis=(0,90), continuous_update = True, wait = False)
#plt.show()

## Let's look at pictographs

### What is a pictograph?

A pictograph is a way to show data using images, where each image represents a certain number of things that are being measured. They look a lot like bar graphs and they can be horizontal or vertical too.

![picture](./images/pictograph-tennis.svg)

from [Math is Fun](https://www.mathsisfun.com/data/images/pictograph-tennis.svg)

### Why do people like to use pictographs?

The main reason is because the pictures offer something for readers to connect with other than just a row of coloured bars.

Also, pictographs often work best to show data with small numbers. If everything can be easily expressed with a simple scale like the one above, then a pictograph might be the right choice to represent the data.

### When are pictographs not a good choice?

In the example above, what if Sam played 46 games instead of 45? This pictogram counts games in steps of 5, so numbers in between these steps might be hard or impossible to show.

A reader might also make a connection with a pictograph that wasn't intended.  Let's show this with an example.

On Halloween, Shayna and Michael went trick-or-treating.  Shayna got 18 pieces of candy, and Michael got 36.  Their totals are shown in this pictograph:

![picture](./images/halloween-candy-collected.jpg)

from [teachersinstitute.yale.edu](http://teachersinstitute.yale.edu/curriculum/units/2008/6/08.06.06.x.html)

At first, is looks like a fine way to show how much candy each child got. The heights of the candy corn pieces are being used to mark the two amounts.  But as a viewer, we don't see just the height—we also see the width. Not only is the second candy corn twice as high, it's also twice as wide, giving it four times the area as the first candy corn.  This makes it *look like* Michael got 4 times as much candy as Shayna, even though he only got twice as much.

Click the "Display" button below to show a better, more accurate way to represent the same data:

hide_me

pic = Image('images/CandyCornGraph.png')
clicker = widgets.Checkbox(value=False, description='Display', disabled=False)

def checking(a):
    if clicker.value == True:
        IPython.display.display(pic)
    else:
        IPython.display.clear_output()
        IPython.display.display(clicker)

IPython.display.display(clicker)
clicker.observe(checking, 'value')

## Let's look at line graphs

### What is a line graph?

A line graph is a way to show how the measurement of one value responds to changes in another, like how something may change over time. In fact, time is one of the most common variables with a line graph.

![picture](./images/line-graph-example.svg)

from [Math is Fun](https://www.mathsisfun.com/data/images/line-graph-example.svg)

### Why are line graphs useful?

They show a moving trend with a line that's easy to follow, instead of just dots on a graph. They work best when the trend involves the change of one variable (jobs, temperature, debt) with respect to another (usually time).

In some cases it can also be useful to plot multiple lines on one graph, usually with different colours to help tell them apart. For example, one might plot polling results for different political parties over time, as with this graph from the CBC:

![federal polling averages](./images/poll-averages.jpg)

from [cbc.ca](http://www.cbc.ca/polopoly_fs/1.3265490!/fileImage/httpImage/image.jpg)

### How can line graphs go wrong?

A common error with line graphs is unlabelled axes. A graph might show a line that slopes upwards, but without labels, we wouldn't know what is growing or how fast. Also, line graphs can trick us into thinking a trend is linear by spacing out the ticks unevenly on one axis, so that the data points neatly line up. Like this example:

![Picture](./images/job-loss-by-quarter.png)

from [Online Stat Book](http://onlinestatbook.com/2/graphing_distributions/graphics/graph2.png)

hide_me

fix = widgets.SelectionSlider(options=['original', 'fixed'], value ='original', description='Slide to fix',
    continuous_update=True, orientation='horizontal',)

def fixing(a):
    if fix.value == 'fixed':
        IPython.display.clear_output()
        IPython.display.display(fix)
        f, ax1 = plt.subplots(1,1,figsize=(10,5))
        ax1.set_title("Job Loss by Quarter")
        ax1.set_xlabel('Months from November 07',fontsize=15)
        ax1.set_ylabel("Jobs Lost in Millions",color='b',fontsize=15)
        x1 = [1, 10, 16, 29]
        y1 = [7,9,13.5,15]
        ax1.plot(x1, y1,"bo-")
        plt.legend()
        plt.show()
    else:
        IPython.display.clear_output()
        IPython.display.display(fix)
        f, ax1 = plt.subplots(1,1,figsize=(10,5))
        ax1.set_title("Job Loss by Quarter")
        ax1.set_xlabel('Months from November 07',fontsize=15)
        ax1.set_ylabel("Jobs Lost in Millions",color='b',fontsize=15)
        x1 = [0,7,23,29]
        y1 = [7,9,13.5,15]
        ax1.plot(x1, y1,"bo-")
        plt.legend()
        plt.show()

IPython.display.display(fix)
fix.observe(fixing, 'value')

IPython.display.clear_output()
IPython.display.display(fix)
f, ax1 = plt.subplots(1,1,figsize=(10,5))
ax1.set_title("Job Loss by Quarter")
ax1.set_xlabel('Months from November 07',fontsize=15)
ax1.set_ylabel("Jobs Lost in Millions",color='b',fontsize=15)
x1 = [0,7,23,29]
y1 = [7,9,13.5,15]
ax1.plot(x1, y1,"bo-")
plt.legend()
plt.show()

## Let's look at circle graphs

### What is a circle graph?

Also known as a pie chart, a circle graph is used to show how a total is split into different groups. The whole pie represents the total, and each slice of the pie represents a different group.  Each slice gets as much of the pie as its group has of the total—the bigger the slice, the more of the total that group represents.

![picture](./images/pie-chart-movies.svg)

from [Math is Fun](https://www.mathsisfun.com/data/images/pie-chart-movies.svg)

### Why are circle graphs useful?

They make it easy to compare group sizes; if there's a biggest or smallest group, that's easy to see, since group sizes are shown as pieces of a whole.

### Why might people not use circle graphs?

To be displayed as a circle graph, data must be converted into percentages of the total, then into slices of a circle, which is more work than other graphs need. Plus, it's easy to mess up if the data are not converted properly (or at all). Circle graphs are also hard to draw accurately on paper, since you need a protractor to ensure your angles are correct. Some people might even say that any time a circle graph would do, a bar graph would do better, and that the pie chart below is the only acceptable one.

![picture](./images/Pie-I-have-Eaten.jpg)

from [Flowing Data](https://i1.wp.com/flowingdata.com/wp-content/uploads/2008/09/Pie-I-have-Eaten.jpg)

********************

### What's wrong with these graphs?

![Picture](./images/unemployment-rate.jpg)

[Business Insider](https://amp.businessinsider.com/images/51cb26c469beddf14c000015-750-453.jpg)

![Picture](./images/candidates-pie.png)

[Flowing Data](http://flowingdata.com/wp-content/uploads/yapb_cache/app15725951258947184.acq6gmp0hf4sowckg80ssc8wg.2xne1totli0w8s8k0o44cs0wc.th.png)

![Picture](./images/130207SuperBowlPoll.jpg)

[Flowing Data](https://i0.wp.com/flowingdata.com/wp-content/uploads/2013/03/130207SuperBowlPoll.jpg?fit=500%2C430&ssl=1)

## What was wrong with these graphs?

1. Mislabeled/Missing axes
2. Plotted wrong
3. Hard to read
4. Numbers don't add to 100%
5. Wrong data shown

The video below goes through several examples of bad/misleading graphs (some of them shown in this notebook) and why they are not good representations of the original data.

hide_me

YouTubeVideo('1F7gm_BG0iQ')

*************

## Practice Questions

### Question 1

A group of kids was asked **what they do first** when they get home from school. The data are shown in the table below. [Data source here](http://www.ur.umich.edu/9900/Apr03_00/7.htm)

| Activity        | Percent|
|-----------------|-----|
|   Eat           | 27% |
|   Personal Care | 19% |
|   Watch TV      | 15% |
|   Study         | 13% |
|   Play          | 9%  |
|   Other         | 17% |

hide_me

answer = widgets.RadioButtons(options=['','circle graph', 'line graph', 'bar graph', 'pictograph'],
                              value='', description='Answer:')
labels = ['Eat', 'Personal Care', 'Watch TV', 'Study', 'Play', 'Other']
data = [0.27, 0.19, 0.15, 0.13, 0.09, 0.17]

def display():
    print('What would be the best graph to display this set of data?')
    IPython.display.display(answer)

def check(a):
    IPython.display.clear_output(wait=False)
    display()
    if answer.value == 'circle graph':
        print("Correct! Circle graphs are used for percentages.")
        print("Let's see this data in a circle graph.")
        patches, texts = plt.pie(data, labels=labels)
        plt.axis('equal')
        plt.tight_layout()
        plt.show()
    else:
        if answer.value == 'bar graph':
            print("A Bar graph would work, but there's a better option. Try again.")
        else:
            if answer.value == 'line graph':
                print("Line graphs are good for change over time, not percentages. Try again.")
            else:
                print("A pictograph would work if the data was in amounts instead of percentages. Try again.")

display()
answer.observe(check, 'value')

### Question 2

A group of kids was asked **how much time** they spend doing different activities after school.  The data are shown in the table below. [Data source here](http://www.ur.umich.edu/9900/Apr03_00/7.htm)

| Activity        | Time spent (minutes)|
|-----------------|-----|
|   Reading       | 30  |
|  Chores         | 30  |
|   Watch TV      | 100 |
|   Study         | 60  |
|   Play          | 74  |
|   Sports        | 60  |

hide_me

answer2 = widgets.RadioButtons(options=['','circle graph', 'line graph', 'bar graph', 'pictograph'],
                              value='', description='Answer:')
labels2 = ['Reading', 'Chores', 'Watch TV', 'Study', 'Play', 'Sports']
data2 = [30, 30, 100, 60, 74, 60]
x = np.arange(len(data2))

def display2():
    print('What would be the best graph to display this set of data?')
    IPython.display.display(answer2)

def check(a):
    IPython.display.clear_output(wait=False)
    display2()
    if answer2.value == 'circle graph':
        print("A circle graph is used for percentages. Try again.")
    else:
        if answer2.value == 'bar graph':
            print("Correct! A bar graph shows the relation between both parameters in an easy to read format.")
            print("Let's see what that looks like.")
            plt.bar(x, data2, width = .3)
            plt.xticks(x, labels2)
            plt.ylabel('Time in Minutes')
            plt.title('Time Spent on Afterschool Activities')
            plt.show()
        else:
            if answer2.value == 'line graph':
                print("Line graphs are good for change over time. Try again.")
            else:
                print("A pictograph would work, but there's a better option to be more accurate. Try again.")

display2()
answer2.observe(check, 'value')

Now that we have seen many examples of both good and bad graphs, let's look at [the worst graphs in science literature](https://www.biostat.wisc.edu/~kbroman/topten_worstgraphs/) and try to figure out why each graph is on the list.

(You can click on a graph to make it bigger.  See if you can point out the flaws in these graphs on your own, then click the "Discussion" links to learn what's wrong with each graph and why it made it to the list!)

*Even top academic institutions sometimes produce images that, while they might look impressive, do very little to help us understand the numbers involved in an issue. Try to figure out what's wrong with these
[infographics from Princeton](http://www.princeton.edu/~ina/infographics/index.html) on your own.*

## What have we learned?

In this notebook, we have learned:

* Graphs are great at conveying numerical information
* Not all types of graphs are created equal
* Graphs can be manipulated to be misleading
* How to identify misleading parts of graphs
* There are steps to create truthful graphs

With this knowledge, you are able to be more informed when you see any kind of data displayed in a graph. You are also able to create truthful graphs of your own. Keep trying to identify good and bad graphs in your every day life.

[![Callysto.ca License](https://github.com/callysto/curriculum-notebooks/blob/master/callysto-notebook-banner-bottom.jpg?raw=true)](https://github.com/callysto/curriculum-notebooks/blob/master/LICENSE.md)


"""Module for compiling codegen output, and wrap the binary for use in
python.

.. note:: To use the autowrap module it must first be imported

   >>> from sympy.utilities.autowrap import autowrap

This module provides a common interface for different external backends, such
as f2py, fwrap, Cython, SWIG(?) etc. (Currently only f2py and Cython are
implemented) The goal is to provide access to compiled binaries of acceptable
performance with a one-button user interface, i.e.

    >>> from sympy.abc import x,y
    >>> expr = ((x - y)**(25)).expand()
    >>> binary_callable = autowrap(expr)
    >>> binary_callable(1, 2)
    -1.0

The callable returned from autowrap() is a binary python function, not a
SymPy object.  If it is desired to use the compiled function in symbolic
expressions, it is better to use binary_function() which returns a SymPy
Function object.  The binary callable is attached as the _imp_ attribute and
invoked when a numerical evaluation is requested with evalf(), or with
lambdify().

    >>> from sympy.utilities.autowrap import binary_function
    >>> f = binary_function('f', expr)
    >>> 2*f(x, y) + y
    y + 2*f(x, y)
    >>> (2*f(x, y) + y).evalf(2, subs={x: 1, y:2})
    0.e-110

The idea is that a SymPy user will primarily be interested in working with
mathematical expressions, and should not have to learn details about wrapping
tools in order to evaluate expressions numerically, even if they are
computationally expensive.

When is this useful?

    1) For computations on large arrays, Python iterations may be too slow,
       and depending on the mathematical expression, it may be difficult to
       exploit the advanced index operations provided by NumPy.

    2) For *really* long expressions that will be called repeatedly, the
       compiled binary should be significantly faster than SymPy's .evalf()

    3) If you are generating code with the codegen utility in order to use
       it in another project, the automatic python wrappers let you test the
       binaries immediately from within SymPy.

    4) To create customized ufuncs for use with numpy arrays.
       See *ufuncify*.

When is this module NOT the best approach?

    1) If you are really concerned about speed or memory optimizations,
       you will probably get better results by working directly with the
       wrapper tools and the low level code.  However, the files generated
       by this utility may provide a useful starting point and reference
       code. Temporary files will be left intact if you supply the keyword
       tempdir="path/to/files/".

    2) If the array computation can be handled easily by numpy, and you
       don't need the binaries for another project.

"""

from __future__ import print_function, division

_doctest_depends_on = { 'exe': ('f2py', 'gfortran'), 'modules': ('numpy',)}

import sys
import os
import shutil
import tempfile
from subprocess import STDOUT, CalledProcessError

from sympy.core.compatibility import check_output
from sympy.utilities.codegen import (
    get_code_generator, Routine, OutputArgument, InOutArgument,
    CodeGenArgumentListError, Result
)
from sympy.utilities.lambdify import implemented_function
from sympy.utilities.decorator import doctest_depends_on
from sympy import C

class CodeWrapError(Exception):
    pass

class CodeWrapper:
    """Base Class for code wrappers"""
    _filename = "wrapped_code"
    _module_basename = "wrapper_module"
    _module_counter = 0

    @property
    def filename(self):
        return "%s_%s" % (self._filename, CodeWrapper._module_counter)

    @property
    def module_name(self):
        return "%s_%s" % (self._module_basename, CodeWrapper._module_counter)

    def __init__(self, generator, filepath=None, flags=[], verbose=False):
        """
        generator -- the code generator to use
        """
        self.generator = generator
        self.filepath = filepath
        self.flags = flags
        self.quiet = not verbose

    @property
    def include_header(self):
        return bool(self.filepath)

    @property
    def include_empty(self):
        return bool(self.filepath)

    def _generate_code(self, main_routine, routines):
        routines.append(main_routine)
        self.generator.write(
            routines, self.filename, True, self.include_header,
            self.include_empty)

    def wrap_code(self, routine, helpers=[]):

        workdir = self.filepath or tempfile.mkdtemp("_sympy_compile")
        if not os.access(workdir, os.F_OK):
            os.mkdir(workdir)
        oldwork = os.getcwd()
        os.chdir(workdir)
        try:
            sys.path.append(workdir)
            self._generate_code(routine, helpers)
            self._prepare_files(routine)
            self._process_files(routine)
            mod = __import__(self.module_name)
        finally:
            sys.path.remove(workdir)
            CodeWrapper._module_counter += 1
            os.chdir(oldwork)
            if not self.filepath:
                shutil.rmtree(workdir)

        return self._get_wrapped_function(mod)

    def _process_files(self, routine):
        command = self.command
        command.extend(self.flags)
        try:
            retoutput = check_output(command, stderr=STDOUT)
        except CalledProcessError as e:
            raise CodeWrapError(
                "Error while executing command: %s. Command output is:\n%s" % (
                    " ".join(command), e.output))
        if not self.quiet:
            print(retoutput)

class DummyWrapper(CodeWrapper):
    """Class used for testing independent of backends """

    template = """# dummy module for testing of SymPy
def %(name)s():
    return "%(expr)s"
%(name)s.args = "%(args)s"
%(name)s.returns = "%(retvals)s"
"""

    def _prepare_files(self, routine):
        return

    def _generate_code(self, routine, helpers):
        with open('%s.py' % self.module_name, 'w') as f:
            printed = ", ".join(
                [str(res.expr) for res in routine.result_variables])
            # convert OutputArguments to return value like f2py
            inargs = filter(lambda x: not isinstance(
                x, OutputArgument), routine.arguments)
            retvals = []
            for val in routine.result_variables:
                if isinstance(val, Result):
                    retvals.append('nameless')
                else:
                    retvals.append(val.result_var)

            print(DummyWrapper.template % {
                'name': routine.name,
                'expr': printed,
                'args': ", ".join([str(arg.name) for arg in inargs]),
                'retvals': ", ".join([str(val) for val in retvals])
            }, end="", file=f)

    def _process_files(self, routine):
        return

    @classmethod
    def _get_wrapped_function(cls, mod):
        return mod.autofunc

class CythonCodeWrapper(CodeWrapper):
    """Wrapper that uses Cython"""

    setup_template = """
from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext

setup(
    cmdclass = {'build_ext': build_ext},
    ext_modules = [Extension(%(args)s)]
        )
"""

    @property
    def command(self):
        command = [sys.executable, "setup.py", "build_ext", "--inplace"]
        return command

    def _prepare_files(self, routine):
        pyxfilename = self.module_name + '.pyx'
        codefilename = "%s.%s" % (self.filename, self.generator.code_extension)

        # pyx
        with open(pyxfilename, 'w') as f:
            self.dump_pyx([routine], f, self.filename,
                self.include_header, self.include_empty)

        # setup.py
        ext_args = [repr(self.module_name), repr([pyxfilename, codefilename])]
        with open('setup.py', 'w') as f:
            print(CythonCodeWrapper.setup_template % {
                'args': ", ".join(ext_args)}, file=f)

    @classmethod
    def _get_wrapped_function(cls, mod):
        return mod.autofunc_c

    def dump_pyx(self, routines, f, prefix, header=True, empty=True):
        """Write a Cython file with python wrappers

           This file contains all the definitions of the routines in c code and
           refers to the header file.

           :Arguments:

           routines
                List of Routine instances
           f
                File-like object to write the file to
           prefix
                The filename prefix, used to refer to the proper header file.
                Only the basename of the prefix is used.
           empty
                Optional. When True, empty lines are included to structure the
                source files. [DEFAULT=True]
        """
        for routine in routines:
            prototype = self.generator.get_prototype(routine)

            # declare
            print('cdef extern from "%s.h":' % prefix, file=f)
            print('   %s' % prototype, file=f)
            if empty:
                print(file=f)

            # wrap
            ret, args_py = self._split_retvals_inargs(routine.arguments)
            args_c = ", ".join([str(a.name) for a in routine.arguments])
            print("def %s_c(%s):" % (routine.name,
                ", ".join(self._declare_arg(arg) for arg in args_py)), file=f)
            for r in ret:
                if not r in args_py:
                    print("   cdef %s" % self._declare_arg(r), file=f)
            rets = ", ".join([str(r.name) for r in ret])
            if routine.results:
                call = '   return %s(%s)' % (routine.name, args_c)
                if rets:
                    print(call + ', ' + rets, file=f)
                else:
                    print(call, file=f)
            else:
                print('   %s(%s)' % (routine.name, args_c), file=f)
                print('   return %s' % rets, file=f)

            if empty:
                print(file=f)
    dump_pyx.extension = "pyx"

    def _split_retvals_inargs(self, args):
        """Determines arguments and return values for python wrapper"""
        py_args = []
        py_returns = []
        for arg in args:
            if isinstance(arg, OutputArgument):
                py_returns.append(arg)
            elif isinstance(arg, InOutArgument):
                py_returns.append(arg)
                py_args.append(arg)
            else:
                py_args.append(arg)
        return py_returns, py_args

    def _declare_arg(self, arg):
        t = arg.get_datatype('c')
        if arg.dimensions:
            return "%s *%s" % (t, str(arg.name))
        else:
            return "%s %s" % (t, str(arg.name))

class F2PyCodeWrapper(CodeWrapper):
    """Wrapper that uses f2py"""

    @property
    def command(self):
        filename = self.filename + '.' + self.generator.code_extension
        command = ["f2py", "-m", self.module_name, "-c", filename]
        return command

    def _prepare_files(self, routine):
        pass

    @classmethod
    def _get_wrapped_function(cls, mod):
        return mod.autofunc

def _get_code_wrapper_class(backend):
    wrappers = { 'F2PY': F2PyCodeWrapper, 'CYTHON': CythonCodeWrapper,
        'DUMMY': DummyWrapper}
    return wrappers[backend.upper()]

@doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
def autowrap(
    expr, language='F95', backend='f2py', tempdir=None, args=None, flags=[],
        verbose=False, helpers=[]):
    """Generates python callable binaries based on the math expression.

    expr
        The SymPy expression that should be wrapped as a binary routine

    :Optional arguments:

    language
        The programming language to use, currently 'C' or 'F95'
    backend
        The wrapper backend to use, currently f2py or Cython
    tempdir
        Path to directory for temporary files.  If this argument is supplied,
        the generated code and the wrapper input files are left intact in the
        specified path.
    args
        Sequence of the formal parameters of the generated code, if ommited the
        function signature is determined by the code generator.
    flags
        Additional option flags that will be passed to the backend
    verbose
        If True, autowrap will not mute the command line backends.  This can be
        helpful for debugging.
    helpers
        Used to define auxillary expressions needed for the main expr.  If the
        main expression need to do call a specialized function it should be put
        in the ``helpers`` list.  Autowrap will then make sure that the compiled
        main expression can link to the helper routine.  Items should be tuples
        with (<funtion_name>, <sympy_expression>, <arguments>).  It is
        mandatory to supply an argument sequence to helper routines.

    >>> from sympy.abc import x, y, z
    >>> from sympy.utilities.autowrap import autowrap
    >>> expr = ((x - y + z)**(13)).expand()
    >>> binary_func = autowrap(expr)
    >>> binary_func(1, 4, 2)
    -1.0

    """

    code_generator = get_code_generator(language, "autowrap")
    CodeWrapperClass = _get_code_wrapper_class(backend)
    code_wrapper = CodeWrapperClass(code_generator, tempdir, flags, verbose)
    try:
        routine = Routine('autofunc', expr, args)
    except CodeGenArgumentListError as e:
        # if all missing arguments are for pure output, we simply attach them
        # at the end and try again, because the wrappers will silently convert
        # them to return values anyway.
        new_args = []
        for missing in e.missing_args:
            if not isinstance(missing, OutputArgument):
                raise
            new_args.append(missing.name)
        routine = Routine('autofunc', expr, args + new_args)

    helps = []
    for name, expr, args in helpers:
        helps.append(Routine(name, expr, args))

    return code_wrapper.wrap_code(routine, helpers=helps)

@doctest_depends_on (exe=('f2py', 'gfortran'), modules=('numpy',))
def binary_function(symfunc, expr, **kwargs):
    """Returns a sympy function with expr as binary implementation

    This is a convenience function that automates the steps needed to
    autowrap the SymPy expression and attaching it to a Function object
    with implemented_function().

    >>> from sympy.abc import x, y
    >>> from sympy.utilities.autowrap import binary_function
    >>> expr = ((x - y)**(25)).expand()
    >>> f = binary_function('f', expr)
    >>> type(f)
    <class 'sympy.core.function.UndefinedFunction'>
    >>> 2*f(x, y)
    2*f(x, y)
    >>> f(x, y).evalf(2, subs={x: 1, y: 2})
    -1.0
    """
    binary = autowrap(expr, **kwargs)
    return implemented_function(symfunc, binary)

@doctest_depends_on (exe=('f2py', 'gfortran'), modules=('numpy',))
def ufuncify(args, expr, **kwargs):
    """
    Generates a binary ufunc-like lambda function for numpy arrays

    ``args``
        Either a Symbol or a tuple of symbols. Specifies the argument sequence
        for the ufunc-like function.

    ``expr``
        A SymPy expression that defines the element wise operation

    ``kwargs``
        Optional keyword arguments are forwarded to autowrap().

    The returned function can only act on one array at a time, as only the
    first argument accept arrays as input.

    .. Note:: a *proper* numpy ufunc is required to support broadcasting, type
       casting and more.  The function returned here, may not qualify for
       numpy's definition of a ufunc.  That why we use the term ufunc-like.

    References
    ==========
    [1] http://docs.scipy.org/doc/numpy/reference/ufuncs.html

    Examples
    ========

    >>> from sympy.utilities.autowrap import ufuncify
    >>> from sympy.abc import x, y
    >>> import numpy as np
    >>> f = ufuncify([x, y], y + x**2)
    >>> f([1, 2, 3], 2)
    [ 3.  6.  11.]
    >>> a = f(np.arange(5), 3)
    >>> isinstance(a, np.ndarray)
    True
    >>> print a
    [ 3. 4. 7. 12. 19.]

    """
    y = C.IndexedBase(C.Dummy('y'))
    x = C.IndexedBase(C.Dummy('x'))
    m = C.Dummy('m', integer=True)
    i = C.Dummy('i', integer=True)
    i = C.Idx(i, m)
    l = C.Lambda(args, expr)
    f = implemented_function('f', l)

    if isinstance(args, C.Symbol):
        args = [args]
    else:
        args = list(args)

    # first argument accepts an array
    args[0] = x[i]
    return autowrap(C.Equality(y[i], f(*args)), **kwargs)