#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 14 16:23:54 2019
@copyright: Copyright 2019, HITS gGmbH
@license: EUPL v. 1.2
@author: richtesn
"""
import os
from shutil import copyfile
import glob
from rpy2.robjects import r
from rpy2.robjects import Vector
# import rpy2's package module
import rpy2.robjects.packages as rpackages
#import rpy2.robjects.numpy2ri
from rpy2.robjects import pandas2ri
import pandas as pd

from rpy2.robjects.packages import importr

R_REPO = "https://cloud.r-project.org"


class StructureAlign:
    """ Class to run pairwise alignments bewtween PDB structure objects.
    """

    def importfromR(self, packname):
        try:
            rpack = rpackages.importr(packname)
        except Exception:
            utils = rpackages.importr("utils")
            # , contriburl = R_REPO)
            utils.install_packages(packname, dependencies=True)
            rpack = rpackages.importr(packname)
        return rpack

    def __init__(self, structures=[], dataDir="."):
        self.__structures = structures
        self.__data_dir = dataDir

    def alignWithTemplate(self, template=None):
        base = self.importfromR('base')
        utils = self.importfromR("utils")
        bio3d = self.importfromR("bio3d")

        for mobiles in glob.glob(self.__data_dir+"/mobile*"):
            self._removeSafe(mobiles)
        for fixed in glob.glob(self.__data_dir+"/fixed"):
            self._removeSafe(fixed)
        for fixed in glob.glob(self.__data_dir+"/*-aln.pdb"):
            self._removeSafe(fixed)
        print("Reading: "+str(os.path.join(self.__data_dir, str(template)+".pdb")))
        template_R = bio3d.read_pdb2(os.path.join(
            self.__data_dir, str(template)+".pdb"))
        inds_t = bio3d.atom_select(
            pdb=template_R, chain=base.c("A", "B"), string="backbone")

        for mstruct in self.__structures:
            print("Reading: "+str(os.path.join(self.__data_dir, str(mstruct)+".pdb")))
            mstruct_R = bio3d.read_pdb2(os.path.join(
                self.__data_dir, str(mstruct)+".pdb"))
            inds_s = bio3d.atom_select(
                pdb=mstruct_R, chain=base.c("A", "B"), string="backbone")
            if template != mstruct:
                aln = bio3d.struct_aln(template_R, mstruct_R, fixed_inds=inds_t, mobile_inds=inds_s,
                                       write_pdbs=True, outpath=self.__data_dir, max_cycles=500, cutoff=0.01)
                inp = [template_R]
                inp.append(mstruct_R)
                alns = bio3d.pdbaln(inp)
                copyfile("aln.fa", mstruct+".aln")
                mobiles = glob.glob(self.__data_dir+"/mobile_*.pdb")
                last = 0
                for m in mobiles:
                    ms = str(m)
                    s = ms.rfind('_')+1
                    e = ms.rfind('.')
                    end = int(ms[s:e])
                    if end > last:
                        last = end
                copyfile(self.__data_dir+"/mobile_"+str(last)+".pdb",
                         self.__data_dir+"/"+mstruct+"-aln.pdb")
                for mobiles in glob.glob(self.__data_dir+"/mobile*"):
                    self._removeSafe(mobiles)
            else:
                copyfile(self.__data_dir+"/"+mstruct+".pdb",
                         self.__data_dir+"/"+mstruct+"-aln.pdb")
            self._removeSafe(self.__data_dir+"/fixed.pdb")

    def _removeSafe(self, file):
        if os.path.exists(file):
            os.remove(file)


if __name__ == "__main__":
    examplePDBs = os.path.join(os.getcwd(), 'data', 'example', 'pdbs')
#    stal=StructureAlign(structures=["AC1", "AC2", "AC3", "AC4", "AC7", "AC8", "AC9", "AC5", "AC6"],
#                        dataDir=examplePDBs)
#    stal.alignWithTemplate(template="AC1")
    stal = StructureAlign(structures=["1hti",  "1b9b", "1ci1", "1i45", "1tim"],
                          dataDir=examplePDBs)
    stal.alignWithTemplate(template="1tim")
    print("Please find aligned structures in "+str(examplePDBs))
