#!/usr/bin/env python

from __future__ import division, print_function

__author__ = "Felix Simkovic"
__date__ = "16 Sep 2017"

import argparse
import Bio.PDB
import Bio.PDB.Polypeptide
import Bio.SeqIO
import math
import multiprocessing
import os
import shutil
import sys
import tempfile
import warnings
warnings.simplefilter("ignore")

SIDECHAIN_TREATMENTS = ["polyala", "allatom"]


def generate_command_line_interface():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-s", dest="side_chain_treatment", nargs="+",
                        default=SIDECHAIN_TREATMENTS,
                        help="Side-chains to be kept")
    parser.add_argument("-n", dest="fragment_number", default=0,
                        type=int, help="Top-N fragments by Flib score [0 means all]")
    parser.add_argument("-t", dest="threads", default=1,
                        type=int, help="Number of threads")
    parser.add_argument("-i", dest="input", nargs="+",
                        required=True, help="Flib fragment library file(s)")
    parser.add_argument("-o", dest="output", required=True,
                        help="Output directory")
    parser.add_argument("-p", dest="database", required=True,
                        help="Path to the PDB structure database")
    return parser.parse_args()


def extract_data_from_flib_line(line):
    line = line.strip().split()
    pdb, chain, start, end, seq = line[:5]
    length, score = line[8], line[-1]
    return pdb.lower(), chain, int(start), int(end), seq, int(length), float(score)


def extract_fragment_coordinates(fragment_entry):
    structure_file, \
        fragment_file, \
        fragment_chain, \
        relative_fragment_start, \
        relative_fragment_end, \
        fragment_sequence, \
        sidechain_treatment \
        = fragment_entry

    structure = read_structure_file(structure_file)
    chain = structure[0][fragment_chain]

    ppb = Bio.PDB.PPBuilder()
    peptides = ppb.build_peptides(chain)
    offset = 0
    for i, pep in enumerate(peptides):
        peptide_sequence = str(pep.get_sequence())
        if fragment_sequence in peptide_sequence:
            offset = relative_fragment_start \
                - peptide_sequence.index(fragment_sequence)
            break

    try:
        absolute_fragment_start = pep[relative_fragment_start - offset].id[1]
        absolute_fragment_end = pep[relative_fragment_end - offset].id[1]
    except:
        print("Error extracting fragment: %s\n" % fragment_file)
        return

    residues_to_keep = range(absolute_fragment_start,
                             absolute_fragment_end + 1)
    for model in structure.copy():
        for chain in model:
            if chain.id == fragment_chain:
                for residue in chain:
                    if int(residue.id[1]) not in residues_to_keep:
                        structure[model.id][chain.id].detach_child(residue.id)
                    elif sidechain_treatment == "polyala":
                        for atom in residue:
                            if atom.name not in ["C", "N", "O", "CA", "CB"]:
                                structure[model.id][chain.id][residue.id].detach_child(
                                    atom.id)

            else:
                structure[model.id].detach_child(chain.id)

    save_fragment_file(fragment_file, structure)
    if not valid_fragment_coords(fragment_file, fragment_sequence):
        print("Fragment file invalid - %s" % fragment_file)
        os.remove(fragment_file)
    else:
        assert os.path.isfile(fragment_file)


def read_structure_file(structure_file):
    parser = Bio.PDB.PDBParser(PERMISSIVE=1)
    return parser.get_structure("", structure_file)


def save_fragment_file(fragment_f, structure):
    s = Bio.PDB.PDBIO()
    s.set_structure(structure)
    s.save(fragment_f)


def valid_fragment_coords(fragment_f, fragment_seq):
    with open(fragment_f, "r") as f_in:
        try:
            sequences = list(Bio.SeqIO.parse(f_in, "pdb-atom"))
        except KeyError:
            return False
        written_seq = str(sequences[0].seq).replace("X", "")
        if len(sequences) > 1:
            return False
        elif written_seq != fragment_seq:
            return False
        return True


def main():
    args = generate_command_line_interface()
    if os.path.isdir(args.output):
        shutil.rmtree(args.output)
    os.mkdir(args.output)
    sys.stdout.write("Provided Flib libraries are %s ...\n" %
                     ", ".join(args.input))

    fragment_entries = []
    for lib_file in args.input:
        tmp_fragment_entries = []
        for fragl in open(lib_file, 'r'):
            if not fragl.strip():
                continue

            fragment_pdb, \
                fragment_chain, \
                fragment_start, \
                fragment_end, \
                fragment_sequence, \
                fragment_length, \
                fragment_score \
                = extract_data_from_flib_line(fragl)

            structure_file = os.path.join(args.database, fragment_pdb + '.pdb')

            for sidechain in args.side_chain_treatment:
                assert sidechain in SIDECHAIN_TREATMENTS
                fragment_name = "{}{}_{}_{}.pdb".format(fragment_pdb,
                                                        fragment_chain,
                                                        fragment_sequence,
                                                        sidechain)
                fragment_file = os.path.join(args.output, fragment_name)
                entry = [structure_file, fragment_file, fragment_chain,
                         fragment_start, fragment_end, fragment_sequence,
                         sidechain]
                tmp_fragment_entries += [(fragment_score, entry)]

        tmp_fragment_entries = sorted(tmp_fragment_entries, key=lambda x: x[0])
        if args.fragment_number > 0:
            fragment_entries.extend(
                tmp_fragment_entries[:args.fragment_number])
        else:
            fragment_entries.extend(tmp_fragment_entries)

    sys.stdout.write("Extracting a total of %d fragments ...\n" %
                     len(fragment_entries))
    pool = multiprocessing.Pool(args.threads)
    fragments_to_extract = [e[1] for e in fragment_entries]
    pool.map(extract_fragment_coordinates, fragments_to_extract)


if __name__ == "__main__":
    main()
