"""
Created on May 23, 2019

@author: Daniel Konecki
"""
import os
import re
import argparse
import numpy as np
import pandas as pd
from tqdm import tqdm
from re import compile
from time import time, sleep
import xml.etree.ElementTree as XMLET
from numpy import floor, triu, nonzero
from multiprocessing import cpu_count, Pool, Lock
from Bio.Alphabet import Gapped
from Bio.Alphabet.IUPAC import IUPACProtein
from Bio import Entrez
from Bio.Seq import Seq
from Bio.Blast import NCBIXML
from Bio.SeqIO import write, parse
from Bio.SeqRecord import SeqRecord
from Bio.PDB.PDBList import PDBList
from Bio.ExPASy import get_sprot_raw
from Bio.Application import ApplicationError
from Bio.Align.Applications import ClustalwCommandline
from Bio.Blast.Applications import NcbiblastpCommandline
from dotenv import find_dotenv, load_dotenv
from SupportingClasses.PDBReference import PDBReference
from SupportingClasses.SeqAlignment import SeqAlignment
from SupportingClasses.AlignmentDistanceCalculator import AlignmentDistanceCalculator
from SupportingClasses.EvolutionaryTraceAlphabet import FullIUPACProtein

try:
    dotenv_path = find_dotenv(raise_error_if_not_found=True)
except IOError:
    dotenv_path = find_dotenv(raise_error_if_not_found=True, usecwd=True)
load_dotenv(dotenv_path)

class DataSetGenerator(object):
    """
    This class is meant to automate data set generation and improve reproducibility for projects which rely on PDB
    structures and alignments generated by BLASTING the proteins represented in those structures.

    Attributes:
        input_path (str/path): The path to the directory where the data for this data set should be stored. It is
        expected to contain at least one directory name ProteinLists which should contain files where each line denotes
        a separate PDB id (four letter code) and the chain of interest (e.g. "1bolA").
        pdb_path (str): The location at which PDB structures should be saved.
        sequence_path (str): The path to a directory where the sequence data can be written in fasta format.
        blast_path (str): The location where BLAST output can be written.
        filtered_blast_path (str): Directory where filtered sequences can be written in fasta format.
        alignment_path (str): The directory to which the initial alignments can be written.
        filtered_alignment_path (str): Path to a directory where the filtered alignment can be written.
        final_alignment_path (str): The directory to which the final alignments can be written.
        protein_data (dict): A dictionary to hold the protein data generated over the course of data set construction.
        It is initially None and is then built up such that it contains the PDB ids parsed in from the protein list
        as keys, referencing dictionaries with each proteins values.
    """

    def __init__(self, input_path):
        """
        Init

        This function overwrites the default init function for the DataSetGenerator class.

        Args:
            input_path (path): The path to the directory where the data for this data set should be stored. It is
            expected to contain at least one directory named ProteinLists which should contain files where each line
            denotes a separate PDB id (four letter code only).
        """
        if input_path is None:
            self.input_path = os.environ.get('INPUT_PATH')
        else:
            self.input_path = input_path
        self.protein_list_path = os.path.join(self.input_path, 'ProteinLists')
        self.pdb_path = os.path.join(self.input_path, 'PDB')
        if not os.path.isdir(self.pdb_path):
            os.makedirs(self.pdb_path)
        self.sequence_path = os.path.join(self.input_path, 'Sequences')
        if not os.path.isdir(self.sequence_path):
            os.makedirs(self.sequence_path)
        self.blast_path = os.path.join(self.input_path, 'BLAST')
        if not os.path.isdir(self.blast_path):
            os.makedirs(self.blast_path)
        self.filtered_blast_path = os.path.join(self.input_path, 'Filtered_BLAST')
        if not os.path.isdir(self.filtered_blast_path):
            os.makedirs(self.filtered_blast_path)
        self.alignment_path = os.path.join(self.input_path, 'Alignments')
        if not os.path.isdir(self.alignment_path):
            os.makedirs(self.alignment_path)
        self.filtered_alignment_path = os.path.join(self.input_path, 'Filtered_Alignments')
        if not os.path.isdir(self.filtered_alignment_path):
            os.makedirs(self.filtered_alignment_path)
        self.final_alignment_path = os.path.join(self.input_path, 'Final_Alignments')
        if not os.path.isdir(self.final_alignment_path):
            os.makedirs(self.final_alignment_path)
        self.protein_data = None

    def identify_protein_sequences(self, data_set_name, protein_list_fn, sources=['PDB'], processes=1, verbose=False):
        """
        Identify Protein Sequences

        This function takes a file with a list of PDB ids and extracts the sequences using the indicated source(s).

        Args:
            data_set_name (str): The name being used for the data set being constructed.
            protein_list_fn (str): The name of the file where the list of PDB ids  can be
            found. Each id is expected to be on its own line and should be formatted as a five letter code (first four
            characters are the PDB ID and the last character is the chain id).
            sources (list): A list of sources in order of preference, sequences will be parsed from sources in that
            order until one is retrieved successfully. Current options are: 'UNP' for Swiss/Uniprot, 'GB' for GenBank,
            and 'PDB' to use the sequence of the PDB being used the the specified chain.
            processes (int): The number of processes to use when performing the BLAST search.
            verbose (bool): Whether to write out information during processing.
        Returns:
        """
        if not os.path.isfile(protein_list_fn):
            raise ValueError('Protein list file not cannot be found at specified location:\n{}'.format(protein_list_fn))
        if verbose:
            print('Importing protein list')
        protein_data = import_protein_list(protein_list_fn=protein_list_fn)
        # Download the PDBs and parse out the query sequences.
        if verbose:
            print('Downloading structures and parsing in query sequences')
            download_pbar = tqdm(total=len(protein_data), unit='Protein IDs')
        seqs_to_write = []
        sequences = {}
        unique_ids = set()

        def update_pdb_download(pdb_seq_data):
            """
            Update PDB Download

            This function serves to update the progress bar for PDB downloading and sequence extraction.

            Args:
                pdb_seq_data (str): The name of the node which has just finished being characterized.
            """
            if pdb_seq_data[1]['Sequence']:
                if pdb_seq_data[1]['Sequence'].seq not in sequences:
                    sequences[pdb_seq_data[1]['Sequence'].seq] = []
                    unique_ids.add(pdb_seq_data[0])
                    seqs_to_write.append(pdb_seq_data[1]['Sequence'])
                sequences[pdb_seq_data[1]['Sequence'].seq].append(pdb_seq_data[0])
            protein_data[pdb_seq_data[0]].update(pdb_seq_data[1])
            if verbose:
                download_pbar.update(1)
                download_pbar.refresh()

        pdb_download_pool = Pool(processes, initializer=init_pdb_processing_pool,
                                 initargs=(self.pdb_path, self.sequence_path, Lock(), sources, verbose))
        for p_id in protein_data:
            pdb_download_pool.apply_async(func=pdb_processing, callback=update_pdb_download,
                                          args=(p_id, protein_data[p_id]['PDB'], protein_data[p_id]['Chain']))
        pdb_download_pool.close()
        pdb_download_pool.join()
        self.protein_data = protein_data
        if verbose:
            print('Unique Sequences Found: {}!'.format(len(sequences)))
        # Write out a fasta file containing all the query sequences for BLASTing
        all_seqs_fn = os.path.join(self.sequence_path, data_set_name + '.fasta')
        if not os.path.isfile(all_seqs_fn):
            with open(all_seqs_fn, 'w') as all_seqs_handle:
                write(seqs_to_write, all_seqs_handle, 'fasta')
        return all_seqs_fn, unique_ids, sequences

    def filter_blast_results(self, unique_ids, e_value_threshold, min_fraction, min_identity, max_identity, blast_fn,
                             processes=1, verbose=False):
        """
        Filter BLAST results

        This function filters BLAST hits for each unique protein of interest in the current data set. Sequences are
        filtered to remove any which have an E value, minimum sequence fraction/coverage, and minimum or maximum
        identity to the query which falls outside the specified bounds. In addition sequences are filtered to only keep
        sequences which consist of the 20 character amino acid alphabet (plus gaps), to remove any sequences which
        include the words: 'artificial', 'fragment', 'low quality', 'partial', or 'synthetic' in their description, and
        any sequences whose taxonomy includes the terms: 'synthetic' or 'artificial'.

        Args:
            unique_ids (set): The protein identifiers in protein_data which represent unique sequences and which can
            be found in the provided blast_fn.
            e_value_threshold (float): The maximum value of e value to consider when filtering returned sequences from
            BLAST.
            min_fraction (float): The minimum percentage (expressed as a decimal) of a query sequence that a BLAST hit
            must cover to be part of the analysis.
            min_identity (float): The minimum percent identity (expressed as a decimal) that a sequence must have to be
            considered as part of the analysis.
            max_identity (float): The maximum percent identity (expressed as a decimal) that a sequence may have to
            still be considered as part of the analysis.
            blast_fn (str): The path to the BLAST result for the current data set.
            processes (int): The number of processes to use when filtering BLAST hits.
            verbose (bool): Whether to write out information during processing.
        """
        for p_id in unique_ids:
            if verbose:
                print(f'Filter Blast Sequences: {p_id}')
            if not os.path.isdir(self.filtered_blast_path):
                os.makedirs(self.filtered_blast_path)
            pileup_fn = os.path.join(self.filtered_blast_path, f'{p_id}.fasta')
            if os.path.isfile(pileup_fn):
                filtered_blast_count = load_filtered_sequences(protein_id=p_id, pileup_fn=pileup_fn,
                                                               min_fraction=min_fraction, min_identity=min_identity,
                                                               max_identity=max_identity)
            else:
                sequences = []
                blast_handle = open(blast_fn, 'r')
                blast_iter = NCBIXML.parse(blast_handle)
                for blast_record in blast_iter:
                    if not blast_record.query.startswith(p_id):
                        continue
                    if verbose:
                        print(f'Filtering BLAST sequences for: {p_id}')
                        filter_blast_pbar = tqdm(total=len(blast_record.alignments), unit='BLAST sequences')

                    def update_filter_blast(filter_blast_seq):
                        """
                        Update Filter Blast

                        This function updates the sequence list with returns from the processing pool and updates the
                        progress bar if verbose output is specified.

                        Args:
                            filter_blast_seq (SeqRecord/None): A sequence for a given BLAST hit which has passed the
                            filter or None if no sequences in for that hit pass.
                        """
                        if filter_blast_seq:
                            sequences.append(filter_blast_seq)
                        if verbose:
                            filter_blast_pbar.update(1)
                            filter_blast_pbar.refresh()

                    blast_filter_pool = Pool(processes, initializer=init_filter_sequences,
                                             initargs=(self.protein_data[p_id]['Sequence'], e_value_threshold,
                                                       min_fraction, min_identity, max_identity, Gapped(IUPACProtein),
                                                       verbose))
                    for alignment in blast_record.alignments:
                        blast_filter_pool.apply_async(func=filter_sequence, callback=update_filter_blast,
                                                      args=(alignment,))
                    blast_filter_pool.close()
                    blast_filter_pool.join()
                blast_handle.close()
                sequences = [self.protein_data[p_id]['Sequence']] + sequences
                with open(pileup_fn, 'w') as pileup_handle:
                    write(sequences=sequences, handle=pileup_handle, format='fasta')
                filtered_blast_count = len(sequences)
            self.protein_data[p_id]['Filtered_BLAST'] = pileup_fn
            self.protein_data[p_id]['Filter_Count'] = filtered_blast_count

    def align_blast_hits(self, unique_ids, msf=True, fasta=True, processes=1, verbose=False):
        """
        Align BLAST Hits

        This function aligns all sequences which were returned by BLAST and passed the specified sequence filters.

        Args:
        unique_ids (set): The protein identifiers in protein_data which represent unique sequences and which can
        be found in the provided blast_fn.
        msf (bool): Whether to write out alignment files in msf format.
        fasta (bool):Whether to write out alignment files in fasta format.
        processes (int): The number of processes to use when aligning sequences for proteins of interest.
        verbose (bool): Whether to write out information during processing.
        """
        if verbose:
            print(f'Aligning BLAST sequences')
            align_blast_pbar = tqdm(total=len(unique_ids), unit='Proteins')
        completed_alignments = {}

        def update_align_blast(aln_res):
            """
            Update Align BLAST

            This function updates the record of which proteins have been processed and what files were produced and
            updates the progress bar if verbose output has been specified.

            Args:
                aln_res (tuple): Element one must be a string corresponding to the protein ID which has been completed,
                elements two and three must be either a string with the path of either the MSF or Fasta (respectively)
                file outputs, or None if that file was not produced.
            """
            completed_alignments[aln_res[0]] = {'MSF_Aln': aln_res[1], 'FA_Aln': aln_res[2]}
            if verbose:
                align_blast_pbar.update(1)
                align_blast_pbar.refresh()

        aln_filter_pool = Pool(processes, initializer=init_align_sequences,
                               initargs=(self.alignment_path, msf, fasta, verbose))
        for p_id in unique_ids:
            aln_filter_pool.apply_async(func=align_sequences, callback=update_align_blast,
                                        args=(p_id, self.protein_data[p_id]['Filtered_BLAST']))
        aln_filter_pool.close()
        aln_filter_pool.join()
        for p_id in completed_alignments:
            self.protein_data[p_id].update(completed_alignments[p_id])

    def identity_filter_alignment(self, unique_ids, max_identity, processes, verbose):
        """

        Args:
            unique_ids (set): The protein identifiers in protein_data which represent unique sequences and which can
            be found in the provided blast_fn.
            max_identity (float): The maximum percent identity (expressed as a decimal) that a sequence may have to
            still be considered as part of the analysis.
            processes (int): The number of processes to use when filtering BLAST hits.
            verbose (bool): Whether to write out information during processing.
        """
        for p_id in unique_ids:
            if verbose:
                print(f'Identity Filter Alignment Sequences: {p_id}')
            if not os.path.isdir(self.filtered_blast_path):
                os.makedirs(self.filtered_blast_path)
            pileup_fn = os.path.join(self.filtered_alignment_path, f'{p_id}.fasta')
            if os.path.isfile(pileup_fn):
                count = 0
                with open(pileup_fn, 'rU') as handle:
                    for _ in parse(handle, 'fasta'):
                        count += 1
            else:
                calculator = AlignmentDistanceCalculator()
                alignment = SeqAlignment(file_name=self.protein_data[p_id]['FA_Aln'], query_id=p_id)
                alignment.import_alignment()
                alignment = alignment.remove_bad_sequences()
                distance_matrix = triu(
                    1.0 - np.array(calculator.get_distance(alignment.alignment, processes=processes)),
                    k=1)
                count = identity_filter(protein_id=p_id, alignment=alignment, distance_matrix=distance_matrix,
                                        identity_filtered_fn=pileup_fn, max_identity=max_identity)
            self.protein_data[p_id]['Final_Count'] = count
            self.protein_data[p_id]['Filtered_Alignment'] = pileup_fn

    def align_identity_filtered(self, unique_ids, msf=True, fasta=True, processes=1, verbose=False):
        """
        Align Identity Filtered

        This function aligns all sequences which passed the identity filtering of.

        Args:
        unique_ids (set): The protein identifiers in protein_data which represent unique sequences and which can
        be found in the provided blast_fn.
        msf (bool): Whether to write out alignment files in msf format.
        fasta (bool):Whether to write out alignment files in fasta format.
        processes (int): The number of processes to use when aligning sequences for proteins of interest.
        verbose (bool): Whether to write out information during processing.
        """
        if verbose:
            print(f'Aligning identity filtered sequences')
            align_identity_pbar = tqdm(total=len(unique_ids), unit='Proteins')
        completed_alignments = {}

        def update_align_identity(aln_res):
            """
            Update Align Identity

            This function updates the record of which proteins have been processed and what files were produced and
            updates the progress bar if verbose output has been specified.

            Args:
                aln_res (tuple): Element one must be a string corresponding to the protein ID which has been completed,
                elements two and three must be either a string with the path of either the MSF or Fasta (respectively)
                file outputs, or None if that file was not produced.
            """
            completed_alignments[aln_res[0]] = {'Final_MSF_Aln': aln_res[1], 'Final_FA_Aln': aln_res[2]}
            if verbose:
                align_identity_pbar.update(1)
                align_identity_pbar.refresh()

        aln_identity_pool = Pool(processes, initializer=init_align_sequences,
                                 initargs=(self.final_alignment_path, msf, fasta, verbose))
        for p_id in unique_ids:
            aln_identity_pool.apply_async(func=align_sequences, callback=update_align_identity,
                                          args=(p_id, self.protein_data[p_id]['Filtered_Alignment']))
        aln_identity_pool.close()
        aln_identity_pool.join()
        for p_id in completed_alignments:
            self.protein_data[p_id].update(completed_alignments[p_id])

    def build_pdb_alignment_dataset(self, protein_list_fn, processes=1, max_target_seqs=2500, e_value_threshold=0.05,
                                    database='customuniref90.fasta', remote=False, min_fraction=0.7, min_identity=0.40,
                                    max_identity=0.98, msf=True, fasta=True, sources=['PDB'], verbose=False):
        """
        Build Dataset

        This method builds a complete data set based on the protein id list specified in the constructor. A summary file
        with statistics about each protein in the data set is written to the input_dir specified for the
        DataSetGenerator instance, named using the data set name from the protein_list_fn specified here.

        Args:
            protein_list_fn (str): The name of the file where the list of PDB ids  can be
            found. Each id is expected to be on its own line and should be formatted as a five letter code (first four
            characters are the PDB ID and the last character is the chain id).
            processes (int): The number of processes to use when performing the BLAST search.
            max_target_seqs (int): The maximum number of hits to look for in the BLAST database.
            e_value_threshold (float): The maximum e-value for a passing hit.
            database (str): The name of the database used to search for paralogs, homologs, and orthologs.
            remote (bool): Whether to perform the call to blastp remotely or not (in this case num_threads is ignored).
            min_fraction (float): The minimum fraction of the query sequence length for a passing hit.
            min_identity (float): The absolute minimum identity for a passing hit.
            max_identity (float): The absolute maximum identity for a passing hit.
            msf (bool): Whether or not to create an msf version of the ClustalW alignment.
            fasta (bool): Whether or not to create an fasta version of the ClustalW alignment.
            sources (list): A list of sources in order of preference, sequences will be parsed from sources in that
            order until one is retrieved successfully. Current options are: 'UNP' for Swiss/Uniprot, 'GB' for GenBank,
            and 'PDB' to use the sequence of the PDB being used the the specified chain.
            verbose (bool): Whether to write out information during processing.
        Return:
            pandas.DataFrame: Summary of sequence counts at the BLAST, filtered BLAST, and filtered alignment stages as
            well as the time it took to generate data for a specific sequence.
        """
        start = time()
        data_set_name = os.path.splitext(os.path.basename(protein_list_fn))[0]
        summary_fn = os.path.join(self.input_path, f'{data_set_name}_summary.tsv')
        if os.path.isfile(summary_fn):
            df = pd.read_csv(summary_fn, header=0, index_col=None, sep='\t')
        else:
            protein_list_fn = os.path.join(self.protein_list_path, protein_list_fn)

            all_seqs_fn, unique_ids, sequences = self.identify_protein_sequences(
                data_set_name=data_set_name, protein_list_fn=protein_list_fn, sources=sources, processes=processes,
                verbose=verbose)
            # BLAST all query sequences at once
            print('BLASTing query sequences')
            blast_fn, hits = blast_query_sequence(protein_id=data_set_name + '_All_Seqs', blast_path=self.blast_path,
                                                  sequence_fn=all_seqs_fn, evalue=e_value_threshold,
                                                  processes=processes, max_target_seqs=max_target_seqs,
                                                  database=database, remote=remote)
            for p_id in hits:
                for related_p_id in sequences[str(self.protein_data[p_id]['Sequence'].seq)]:
                    self.protein_data[related_p_id].update(hits[p_id])
                    self.protein_data[related_p_id]['BLAST'] = blast_fn
            # Filter the BLAST hits for each query
            print('Filtering BLAST hits')
            self.filter_blast_results(unique_ids, e_value_threshold, min_fraction, min_identity, max_identity, blast_fn,
                                      processes, verbose)
            # Aligning BLAST hits for each query
            print('Aligning BLAST hits')
            self.align_blast_hits(unique_ids, msf, fasta, processes, verbose)
            # Performing identity filter on aligned sequences
            print('Performing Identity filter')
            self.identity_filter_alignment(unique_ids, max_identity, processes, verbose)
            # Aligning identity filter sequences
            print('Aligning Identity filtered sequences')
            self.align_identity_filtered(unique_ids, msf, fasta, processes, verbose)

            summary = {'Protein_ID': [], 'Chain': [], 'Accession': [], 'Sequence_Length': [], 'BLAST_Hits': [],
                       'Filtered_BLAST': [], 'Filtered_Alignment': [], 'Total_Size': []}
            for p_id in unique_ids:
                for related_p_id in sequences[str(self.protein_data[p_id]['Sequence'].seq)]:
                    summary['Protein_ID'].append(related_p_id)
                    summary['Chain'].append(len(str(self.protein_data[p_id]['Chain'])))
                    summary['Accession'].append(len(str(self.protein_data[p_id]['Accession'])))
                    summary['Sequence_Length'].append(len(str(self.protein_data[p_id]['Sequence'])))
                    summary['BLAST_Hits'].append(self.protein_data[p_id]['BLAST_Hits'])
                    summary['Filtered_BLAST'].append(self.protein_data[p_id]['Filter_Count'])
                    summary['Filtered_Alignment'].append(self.protein_data[p_id]['Final_Count'])
                    summary['Total_Size'].append(len(str(self.protein_data[p_id]['Sequence'])) *
                                                 self.protein_data[p_id]['Final_Count'])
            df = pd.DataFrame(summary)
            df.to_csv(summary_fn, sep='\t', index=False, header=True,
                      columns=['Protein_ID', 'Chain', 'Accession', 'Sequence_Length', 'BLAST_Hits', 'Filtered_BLAST',
                               'Filtered_Alignment', 'Total_Size'])
        end = time()
        print(f'It took {(end - start) / 60.} min to generate data set')
        return df

def import_protein_list(protein_list_fn, verbose=False):
    """
    Import protein list

    This function opens the list file found at protein_list_fn and parses in the PDB ids and the chain of interest
    specified there.

    Args:
        protein_list_fn (str): The name of the file where the list of PDB ids  can be found. Each id is expected to be
        on its own line and should be formatted as a five letter code (first four characters are the PDB ID and the last
        character is the chain id).
        verbose (bool): Whether or not to print out updates about the import.
    Returns:
        dict: A dictionary where each key is a single PDB id parsed from the specified list file and each value is
        a dictionary with the first key and value being the specified chain of interest, and will be filled as the
        data set is constructed.
    """
    if verbose:
        print(f'Importing protein list from: {protein_list_fn}')
    protein_list_fn = os.path.join(protein_list_fn)
    protein_list = {}
    pdb_id_pattern = compile(r'^([0-9][a-z0-9]{3})([A-Z])$')
    with open(protein_list_fn, mode='r') as protein_list_handle:
        for line in protein_list_handle:
            pdb_id_match = pdb_id_pattern.match(line.strip())
            try:
                protein_list[pdb_id_match.group(0)] = {'PDB': pdb_id_match.group(1), 'Chain': pdb_id_match.group(2)}
            except AttributeError:
                raise ValueError(f"Encountered bad formatting in protein list: '{line}'")
    if verbose:
        print(f'Imported {len(protein_list)} protein IDs')
    return protein_list

def download_pdb(pdb_path, pdb_id, verbose=False):
    """
    Download PDB

    This function downloads the PDB structure file for the given PDB id provided. The file is stored at a path within a
    sub directory of the provided pdb_path named <pdb_path>/{middle two characters of the PDB id}/pdb{pdb id}.ent. If
    the PDB structure is not available then an attempt is made to download it as an 'obsolete' PDB structure. If this is
    also unsuccessful, None is returned.

    Args:
        pdb_path (str): The location at which PDB structures should be saved.
        pdb_id (str): Four letter code for a PDB id to be downloaded.
        verbose (bool): Whether to write out information during processing.
    Returns:
        str: The path to the PDB file downloaded.
    """
    if verbose:
        print(f'Downloading structure data for: {pdb_id}')
    if pdb_id is None:
        raise ValueError('None is not a valid value for pdb_id!')
    if not os.path.isdir(pdb_path):
        if verbose:
            print(f'Making dir: {pdb_path}')
        os.makedirs(pdb_path)
    pdb_list = PDBList(server='ftp://ftp.wwpdb.org', pdb=pdb_path, verbose=verbose)
    pdb_file = pdb_list.retrieve_pdb_file(pdb_code=pdb_id, file_format='pdb')
    if not os.path.isfile(pdb_file):
        pdb_file = pdb_list.retrieve_pdb_file(pdb_code=pdb_id, file_format='pdb', obsolete=True)
        if not os.path.isfile(pdb_file):
            pdb_file = None
            print(f'Structure does not exist in obsolete: {pdb_id}')
    if verbose:
        print(f'Completed {pdb_id} structure download')
    return pdb_file

def parse_query_sequence(protein_id, pdb_id, chain_id, sequence_path, pdb_fn, sources, verbose=False):
    """
    Parse Query Sequence

    This function opens a downloaded PDB file for a given protein id (for which download_pdb has already been
    called) and extracts the sequence of the specified chain_id for the structure. If Swiss/UniProt or GenBank data is
    available in the structure file header an attempt is made to retrieve the full sequence referenced by the specified
    chain in the structure. If Swiss/UniProt or GenBank information is not available the sequence parsed from the
    structure file is given in single letter amino acid codes. The sequence is saved to a file in sequence_path and is
    given a file name with the pattern {protein id}.fasta. If the specified chain_id is not available for the structure
    the first chain (alphabetically) is used. The chain is returned as the last value so the user knows if the specified
    chain was used or not.

    Args:
        protein_id (str): Full protein ID code (five letter) code for which the sequence should be parsed.
        pdb_id (str): Four letter code for a PDB id whose sequence should be parsed.
        chain_id (str/char): A single letter code for the chain to be extracted.
        sequence_path (str): The path to a directory where the sequence data can be written in fasta format.
        pdb_fn (str): The full path to the PDB from which the sequence should be extracted.
        sources (list): A list of sources in order of preference, sequences will be parsed from sources in that order
        until one is retrieved successfully. Current options are: 'UNP' for Swiss/Uniprot, 'GB' for GenBank, and 'PDB'
        to use the sequence of the PDB being used the the specified chain.
        verbose (bool): Whether or not to write out warnings from PDB parsing.
    Returns:
        str: The sequence parsed from the PDB file of the specified protein id.
        int: The length of the parsed sequence.
        str: The file path to the fasta file where the sequence has been written.
        str: The chain id for which sequence was parsed.
        str: The accession id for the protein data which was parsed (if taken from an external source, either
        Swiss/Uniprot or GenBank).
    """
    if verbose:
        print(f'Parsing query sequence for: {protein_id}')
    if not os.path.isdir(sequence_path):
        if verbose:
            print(f'Making dir: {sequence_path}')
        os.makedirs(sequence_path)
    protein_fasta_fn = os.path.join(sequence_path, f'{protein_id}.fasta')
    if os.path.isfile(protein_fasta_fn):
        if verbose:
            print(f'Loading query sequence from file: {protein_fasta_fn}')
        with open(protein_fasta_fn, 'r') as protein_fasta_handle:
            seq_iter = parse(handle=protein_fasta_handle, format='fasta')
            sequence = next(seq_iter)
            seq_len = len(sequence)
            sequence.alphabet = FullIUPACProtein()
            match = re.match(r'.*Target Query: Chain: ([A-Z])(: From ((UniProt)|(GenBank)) Accession: (.*))?$',
                             sequence.description)
            chain = match.group(1)
            try:
                accession = match.group(6)
            except IndexError:
                accession = None
    else:
        curr_pdb = PDBReference(pdb_file=pdb_fn)
        curr_pdb.import_pdb(structure_id=pdb_id)
        accession = None
        seq = None
        chain = None
        source = None
        remaining_chains = list(sorted(set(curr_pdb.chains) - {chain_id}))
        for chain in [chain_id] + remaining_chains:
            for source in sources:
                try:
                    accession, seq = curr_pdb.get_sequence(chain=chain_id, source=source)
                except ValueError:
                    accession = 'INSPECT MANUALLY'
                if seq or (accession == 'INSPECT MANUALLY'):
                    break
            if seq or (accession == 'INSPECT MANUALLY'):
                break
        if seq:
            desc = (f'Target Query: Chain: {chain}' +
                    (f': From UniProt Accession: {accession}' if source == 'UNP' else '') +
                    (f': From GenBank Accession: {accession}' if source == 'GB' else ''))
            sequence = SeqRecord(Seq(seq, alphabet=FullIUPACProtein()), id=protein_id, description=desc)
            seq_len = len(sequence)
            seq_records = [sequence]
            with open(protein_fasta_fn, 'w') as protein_fasta_handle:
                write(sequences=seq_records, handle=protein_fasta_handle, format='fasta')
        else:
            sequence = None
            protein_fasta_fn = None
            seq_len = 0
            chain = chain_id
    if verbose:
        print(f'Parsed query sequence for: {protein_id}')
    return sequence, seq_len, protein_fasta_fn, chain, accession

def init_pdb_processing_pool(pdb_path, sequence_path, lock, sources, verbose):
    """
    Init PDB Processing Pool

    This function initializes variables needed by all processes downloading PDBs and processing their query sequence.

    Args:
        pdb_path (str): The location at which PDB structures should be saved.
        sequence_path (str): The path to a directory where the sequence data can be written in fasta format.
        lock (multiprocessing.Lock): A lock to control filesystem access while downloading PDBs since many folders need
        to be checked for and created.
        sources (list): A list of sources in order of preference, sequences will be parsed from sources in that order
        until one is retrieved successfully. Current options are: 'UNP' for Swiss/Uniprot, 'GB' for GenBank, and 'PDB'
        to use the sequence of the PDB being used the the specified chain.
        verbose (bool): Whether or not to write out verbose output.
    """
    global pdb_dir, sequence_dir, fs_lock, allowed_src, verbose_out
    pdb_dir = pdb_path
    sequence_dir = sequence_path
    fs_lock = lock
    allowed_src = sources
    verbose_out = verbose

def pdb_processing(protein_id, pdb_id, chain_id):
    """
    PDB Processing

    This function serves to download a single PDB id's structure and parse out its query sequence.

    Args:
        protein_id (str): The full unique protein ID for the protein being processed.
        pdb_id (str): The four letter code for the protein being processed.
        chain_id (str): The single letter code for the chain to be parsed from the provided PDB.
    Returns:
        str: The protein_id passed in, needed for indexing upon return
        dict: A dictionary containing data to be added to the protein_data field of the DataSetGenerator including the
        following keys and values:
            PDB_FN (str): The full path to the PDB structure downloaded for this protein.
            Chain (str): The chain for which the sequence was parsed, should be the same as that in the protein list
            file, but may change if that chain is not available. In that case the first (alphabetical) chain is used.
            Sequence (Bio.SeqRecord.SeqRecord): A SeqRecord containing the protein_id and the parsed out sequence.
            Length (int): The length of the parsed sequence.
            Seq_Fasta (str): The full path to the fasta file containing just the query sequence for this protein.
    """
    fs_lock.acquire()
    pdb_fn = download_pdb(pdb_path=pdb_dir, pdb_id=pdb_id,  verbose=verbose_out)
    fs_lock.release()
    if pdb_fn is None:
        seq, length, seq_fn, chain_id, ext_id = None, None, None, None, None
    else:
        seq, length, seq_fn, chain_id, ext_id = parse_query_sequence(protein_id=protein_id, pdb_id=pdb_id,
                                                                     chain_id=chain_id, sequence_path=sequence_dir,
                                                                     pdb_fn=pdb_fn, sources=allowed_src,
                                                                     verbose=verbose_out)
    data = {'PDB_FN': pdb_fn, 'Chain': chain_id, 'Sequence': seq, 'Length': length, 'Seq_Fasta': seq_fn,
            'Accession': ext_id}
    return protein_id, data

def blast_query_sequence(protein_id, blast_path, sequence_fn, evalue=0.05, processes=1, max_target_seqs=20000,
                         database='nr', remote=False, verbose=False):
    """
    BlAST Query Sequence

    This function uses a local instance of the BLAST tool and the specified database to search for homologs and
    orthologs of the specified protein. The BLAST results are stored in a subdirectory of the input_path named BLAST
    with a file name following the pattern {protein id}.xml. This method assumes that parse_query_sequence has
    already been performed for the specified protein id.

    Args:
        protein_id (str): Four letter code for the PDB id whose sequence should be searched using BLAST.
        blast_path (str): The location where BLAST output can be written.
        sequence_fn (str): The full path to a fasta formatted file which can be used as input for BLAST.
        evalue (float): The E-value upper limit for BLAST hits
        processes (int): The number of threads to use when performing the BLAST search.
        max_target_seqs (int): The maximum number of hits to look for in the BLAST database.
        database (str): The name of the database used to search for paralogs, homologs, and orthologs.
        remote (bool): Whether to perform the call to blastp remotely or not (in this case num_threads is ignored).
        verbose (bool): Whether to write out information during processing.
    Returns:
        str: The path to the xml file storing the BLAST output.
        dict: Mapping between protein_ids and a dictionary of the BLAST_Hit key and the number of hits returned.
    """
    if verbose:
        print('BLASTing query sequence for: {}'.format(protein_id))
    if not os.path.isdir(blast_path):
        if verbose:
            print('Making dir: {}'.format(blast_path))
        os.makedirs(blast_path)
    blast_fn = os.path.join(blast_path, '{}.xml'.format(protein_id))
    if not os.path.isfile(blast_fn):
        start = time()
        if remote:
            blastp_cline = NcbiblastpCommandline(cmd=os.path.join(os.environ.get('BLAST_PATH'), 'blastp'),
                                                 out=blast_fn, query=sequence_fn, outfmt=5, remote=True, ungapped=False,
                                                 max_target_seqs=max_target_seqs, evalue=evalue, db=database)
        else:
            blastp_cline = NcbiblastpCommandline(cmd=os.path.join(os.environ.get('BLAST_PATH'), 'blastp'),
                                                 out=blast_fn, query=sequence_fn, outfmt=5, remote=False,
                                                 ungapped=False, num_threads=processes, evalue=evalue,
                                                 max_target_seqs=max_target_seqs,
                                                 db=os.path.join(os.environ.get('BLAST_DB_PATH'), database))
        if verbose:
            print(blastp_cline)
        print(blastp_cline)
        stdout, stderr = blastp_cline()
        if verbose:
            print(stdout)
            print(stderr)
        end = time()
    else:
        start = end = time()
        if verbose:
            print('BLAST previously completed for: {}'.format(protein_id))
    hit_counts = {}
    with open(blast_fn, 'r') as blast_handle:
        blast_iter = NCBIXML.parse(blast_handle)
        for blast_record in blast_iter:
            query_id = blast_record.query.split()[0]
            hit_counts[query_id] = {'BLAST_Hits': len(blast_record.alignments)}
    if verbose:
        print('BLAST for {} completed in {} mins'.format(protein_id, (end - start) / 60.0))
    return blast_fn, hit_counts

def load_filtered_sequences(protein_id, pileup_fn, min_fraction, min_identity, max_identity):
    """
    Load Filtered Sequences

    If the BLAST hits for a given protein have already been filtered this method checks that all sequences being loaded
    (possibly from a previous run) pass the current sequence filters and counts the sequences.

    Args:
        protein_id (str): The protein whose file is being processed.
        pileup_fn (str): A fasta file containing unaligned sequences which have been parsed from a set of BLAST hits,
        and which should pass the specified filters.
        min_fraction (float): The minimum percentage (expressed as a decimal) of a query sequence that a BLAST hit
        must cover to be part of the analysis.
        min_identity (float): The minimum percent identity (expressed as a decimal) that a sequence must have to be
        considered as part of the analysis.
        max_identity (float): The maximum percent identity (expressed as a decimal) that a sequence may have to
        still be considered as part of the analysis.
    Return:
        int: The number of sequences in the file, if all sequences pass the provided filters.
    """
    hsp_data_pattern = compile(r'^.*\sHSP_identity=(\d+)\sHSP_alignment_length=(\d+)\sFraction_length=([0-9]+[.][0-9]+)\sHSP_taxonomy=([a-zA-Z;]*).*$')
    sequences = []
    with open(pileup_fn, 'r') as pileup_handle:
        fasta_iter = parse(handle=pileup_handle, format='fasta')
        for seq_record in fasta_iter:
            if seq_record.description.startswith(protein_id):  # Add the target sequence without checking.
                sequences.append(seq_record)
                continue
            hsp_data_match = hsp_data_pattern.match(seq_record.description)
            subject_fraction = float(hsp_data_match.group(3))
            if subject_fraction < min_fraction:
                raise ValueError('Sequences in the filtered fasta do not meet the length fraction requirement.\n'
                                 '{}: Fraction={}'.format(seq_record.id, subject_fraction))
            seq_id = int(hsp_data_match.group(1))
            seq_len = int(hsp_data_match.group(2))
            identity = seq_id / float(seq_len)
            if identity < min_identity or identity > max_identity:
                raise ValueError('Sequences in the filtered fasta do not met the identity requirement.\n'
                                 '{}: Identity={}'.format(seq_record.id, identity))
            seq_record.alphabet = FullIUPACProtein()
            sequences.append(seq_record)
    return len(sequences)

def init_filter_sequences(query_seq, e_value_threshold, min_fraction, min_identity, max_identity,
                          alphabet=Gapped(IUPACProtein), verbose=False):
    """
    Init Filter Sequences

    This function initializes a pool of processes which can be used to filter sequences from a given BLAST search.

    Args:
        query_seq: The sequence for the protein being queried.
        e_value_threshold (float): The maximum value of e value to consider when filtering returned sequences from
        BLAST.
        min_fraction (float): The minimum percentage (expressed as a decimal) of a query sequence that a BLAST hit
        must cover to be part of the analysis.
        min_identity (float): The minimum percent identity (expressed as a decimal) that a sequence must have to be
        considered as part of the analysis.
        max_identity (float): The maximum percent identity (expressed as a decimal) that a sequence may have to
        still be considered as part of the analysis.
        alphabet (Bio.Alphabet): An alphabet object specifying the allowed characters for filtered sequences.
        verbose (bool): Whether to write out information during processing.
    """
    global bf_query_seq, bf_e_thresh, bf_min_frac, bf_min_id, bf_max_id, bf_alpha, bf_verbose
    bf_query_seq = query_seq
    bf_e_thresh = e_value_threshold
    bf_min_frac = min_fraction
    bf_min_id = min_identity
    bf_max_id = max_identity
    bf_alpha = alphabet
    bf_verbose = verbose

def filter_sequence(blast_alignment):
    """
    Filter Sequence

    This function parses the HSPs for a given BLAST hit, returning a sequence if any pass the specified sequence filters
    (values provided by init_filter_sequences). If more than one HSP for a given hit passes, the one with the highest
    identity is returned.

    Args:
        blast_alignment (Bio.BLAST.Record.alignment): One alignment in a given BLAST record, the HSPs for which should
        be checked to see if they pass all filters.
    Return:
        SeqRecord: A sequence which passes the specified filters, None if no HSP passes.
    """
    aln_seq_record = None
    aln_identity = None
    for hsp in blast_alignment.hsps:
        if hsp.expect <= bf_e_thresh:  # Should already be controlled by BLAST e-value
            subject_length = len(hsp.sbjct.replace('-', ''))
            query_length = len(bf_query_seq.seq)
            try:
                subject_fraction = min(subject_length / float(query_length), query_length / float(subject_length))
            except ZeroDivisionError:
                continue
            if bf_min_frac > subject_fraction:
                continue
            identity = hsp.identities / float(hsp.align_length)
            if bf_min_id > identity or identity > bf_max_id:
                continue
            if aln_identity and identity <= aln_identity:
                continue
            if 0 in [aa in bf_alpha.letters for aa in set(hsp.sbjct)]:
                continue
            if any([c in blast_alignment.hit_def.lower() for c in
                    {'artificial', 'fragment', 'low quality', 'partial', 'synthetic'}]):
                continue
            try:
                lineage = ';'.join(parse_lineage(blast_alignment.hit_id)).lower()
            except AttributeError:
                if bf_verbose:
                    print(f'No Lineage: {blast_alignment.hit_id}')
                lineage = ''
            if any([c in lineage for c in {'synthetic', 'artificial'}]):
                continue
            new_description = f'{blast_alignment.hit_def} HSP_identity={hsp.identities} HSP_alignment_length='\
                              f'{hsp.align_length} Fraction_length={subject_fraction} HSP_taxonomy={lineage}'
            aln_seq_record = SeqRecord(
                Seq(hsp.sbjct, alphabet=FullIUPACProtein()),
                id=blast_alignment.hit_id, name=blast_alignment.title,
                description=new_description)
            aln_identity = identity
    return aln_seq_record

def init_align_sequences(alignment_path, msf=True, fasta=True, verbose=False):
    """
    Init Align Sequences

    This function initializes a multi processing pool for aligning sequences.

    Args:
        alignment_path (str): The path to the directory where alignment files should be written.
        msf (bool): Whether to write out alignment files in msf format.
        fasta (bool):Whether to write out alignment files in fasta format.
        verbose (bool): Whether to write out information during processing.
    """
    global pool_aln_path, pool_msf, pool_fasta, pool_verbose
    pool_aln_path = alignment_path
    pool_msf = msf
    pool_fasta = fasta
    pool_verbose = verbose

def align_sequences(protein_id, pileup_fn):
    """
    Align Sequences

    This method uses ClustalW to align the query sequence and all of the hits from BLAST which passed the filtering
    process by default the alignment is performed once, to produce a fasta alignment file, and then converted to
    produce the msf alignment file, however either of these options can be turned off.

    Args:
        protein_id (str): The unique protein identifier for which an alignment should be performed.
        pileup_fn (str): The full path to the file with the sequences which should be aligned.
    Returns:
        str: The protein ID for the alignment which has been completed.
        str: The path to the fasta alignment produced by this method (None if fa=False).
        str: The path to the msf alignment produced by this method (None if msf=False).
    """
    print(f'Align sequences: {protein_id}')
    clustalw_path = os.environ.get('CLUSTALW_PATH')
    if not os.path.isdir(pool_aln_path):
        os.makedirs(pool_aln_path)
    if protein_id is None:
        raise ValueError('The protein ID is None, the protein ID is required for file naming and job tracking.')
    if pileup_fn is None:
        raise ValueError('Pileup filename provided is None, pileup file is required to perform an alignment.')
    fa_fn = None
    if pool_fasta:
        fa_fn = os.path.join(pool_aln_path, f'{protein_id}.fasta')
        if not os.path.isfile(fa_fn):
            fa_cline = ClustalwCommandline(clustalw_path, infile=pileup_fn, align=True, quicktree=True, outfile=fa_fn,
                                           output='FASTA')
            if pool_verbose:
                print(fa_cline)
            try:
                stdout, stderr = fa_cline()
                if pool_verbose:
                    print(stdout)
                    print(stderr)
            except ApplicationError:
                fa_fn = None
    msf_fn = None
    if pool_msf:
        msf_fn = os.path.join(pool_aln_path, f'{protein_id}.msf')
        if not os.path.isfile(msf_fn):
            if pool_fasta:
                msf_cline = ClustalwCommandline(clustalw_path, infile=fa_fn, convert=True, outfile=msf_fn, output='GCG')
            else:
                msf_cline = ClustalwCommandline(clustalw_path, infile=pileup_fn, align=True, quicktree=True,
                                                outfile=msf_fn, output='GCG')
            if pool_verbose:
                print(msf_cline)
            try:
                stdout, stderr = msf_cline()
                if pool_verbose:
                    print(stdout)
                    print(stderr)
            except ApplicationError:
                msf_fn = None
    return protein_id, msf_fn, fa_fn

def identity_filter(protein_id, alignment, distance_matrix, identity_filtered_fn, max_identity=0.98):
    """
    Identity Filter

    This method accepts an alignment and distance matrix for a previously generated alignment and then filters the
    alignment so that no sequences have an identity greater than max_identity. If two or more sequences have an identity
    greater than max_identity the first sequence in that group is kept and all others are discarded (unless one of the
    sequences is the query sequence, in which case the query sequence is kept and all other sequences in the group are
    discarded).

    Args:
        protein_id (str): The PDB id for the protein whose alignment is being filtered.
        alignment (SeqAlignment): The alignment which has been previously generated.
        distance_matrix (np.array): The distance matrix for the passed in alignment.
        identity_filtered_fn (str): The path to which to write the filtered sequences.
        max_identity (float): The maximum sequence identity two sequences can share.
    Returns:
        int: The number of sequences which pass through this filtering process.
    """
    to_keep = set()
    to_remove = set()
    query_seq_pos = alignment.seq_order.index(protein_id)
    to_keep.add(query_seq_pos)
    for i in range(alignment.size):
        row_ids = distance_matrix[i, :]
        above_max_id = row_ids > max_identity
        positions_to_remove = set(nonzero(above_max_id)[0])
        if i in to_keep or i in to_remove:
            to_remove.update(positions_to_remove)
        elif positions_to_remove.isdisjoint(to_keep):
            to_keep.add(i)
            to_remove.update(positions_to_remove)
        else:
            to_remove.add(i)
    filtering_count = (len(to_keep) + len(to_remove))
    if alignment.size != filtering_count:
        raise ValueError('Identity filtering does not match alignment size {} != {} = {} + {}'.format(
            alignment.size, filtering_count, len(to_keep), len(to_remove)))
    filtered_alignment = alignment.generate_sub_alignment([alignment.seq_order[x] for x in to_keep])
    filtered_sequences = []
    for record in filtered_alignment.alignment:
        record.seq = record.seq.ungap(gap='-')
        filtered_sequences.append(record)
    with open(identity_filtered_fn, 'w') as out_handle:
        write(filtered_sequences, out_handle, 'fasta')
    count = filtered_alignment.size
    return count

def parse_genbank_lineage(full_id):
    """
    Parse GenBank Lineage

    This method retrieves the information for a specific GenBank entry (if possible) and extracts the taxonomy
    information from that entry.

    Args:
        full_id (str): A GenBank identifier for which to look up the taxonomy on the NCBI servers.
    Return:
        list: An ordered listing of the provided identifier's taxonomy going from most coarse to most specific.
    """
    acc_pattern = re.compile(r'^(dbj|emb|gb|pdb|ref|sp)\|(.*)\|([A-Z])?$')
    match_groups = acc_pattern.match(full_id)
    if match_groups.group(1) == 'pdb':
        acc = f'{match_groups.group(2)}_{match_groups.group(3)}'
    else:
        acc = match_groups.group(2)
    Entrez.email = os.environ.get('EMAIL')
    handle = None
    attempts = 0
    success = False
    while (not success) and (attempts < 10):
        try:
            handle = Entrez.efetch(db='protein', rettype='gp', retmode='xml', id=acc)
            success = True
            print('Success')
        except AttributeError as e:
            attempts += 1
            sleep(1)
            print('FAILED')
            # print(full_id)
            # print(acc)
            # raise e
    if handle is None:
        raise ValueError(f'{full_id} with acc: {acc} could not be retrieved!')
    gb_set = XMLET.parse(handle).getroot()
    gb_seq = gb_set[0]
    gb_taxa = gb_seq.find('GBSeq_taxonomy')
    full_lineage = gb_taxa.text.replace(' ', '').split(';')
    return full_lineage

def parse_uniref_lineage(full_id):
    """
    Parse UniRef Lineage

    This method retrieves the information for a specific UniRef entry (if possible) and extracts the taxonomy
    information from that entry.

    Args:
        full_id (str): A UniRef identifier for which to look up the taxonomy on the UniRef/UniProt servers.
    Return:
        list: An ordered listing of the provided identifier's taxonomy going from most coarse to most specific.
    """
    acc_pattern = re.compile(r'^UniRef\d+_([0-9A-Z]+)$')
    acc = acc_pattern.match(full_id).group(1)
    handle = get_sprot_raw(acc)
    data = handle.read()
    lines = data.split('\n')
    taxonomy_pattern = re.compile(r'^OC\s+(.*)[;\.]$')
    taxonomy_lines = [taxonomy_pattern.match(line).group(1) for line in lines if line.startswith('OC')]
    full_lineage = []
    for tax_line in taxonomy_lines:
        full_lineage += tax_line.split('; ')
    return full_lineage

def parse_lineage(full_id):
    """
    Parse ULineage

    This method retrieves the information for a specific identifier (if possible), trying first the UniRef/UniProt
    database and then the GenBank/NCBI database. The taxonomy information is extracted for that entry.

    Args:
        full_id (str): A UniRef or GenBank sequence identifier for which to look up the taxonomy.
    Return:
        list: An ordered listing of the provided identifier's taxonomy going from most coarse to most specific.
    """
    try:
        full_lineage = parse_uniref_lineage(full_id)
    except Exception as e:
        try:
            full_lineage = parse_genbank_lineage(full_id)
        except Exception as e:
            raise e
    return full_lineage

def determine_identity_bin(identity_count, length, interval, abs_max_identity, abs_min_identity, identity_bins):
    """
    Determine Identity Bin

    This method determines which identity bin a sequence belongs in based on the settings used for filtering
    sequences in the filter_blast_sequences function.

    Args:
        identity_count (int): The number of positions which match between the query and the BLAST hit.
        length (int): The number of positions in the aligned sequence (including gaps).
        interval (int): The interval on which to define bins between min_identity and abs_min_identity in case
        sufficient sequences are not found at min_identity.
        abs_max_identity (float): The absolute maximum identity for a passing hit.
        abs_min_identity (float): The absolute minimum identity for a passing hit.
        identity_bins (set): All of the possible identity bin options.
    Returns:
        float: The identity bin in which the sequence belongs.
    """
    similarity = (identity_count / float(length)) * 100
    similarity_int = floor(similarity)
    similarity_bin = similarity_int - (similarity_int % (interval * 100))
    similarity_bin /= 100
    final_bin = None
    # If the similarity_bin is within the specified range of identities.
    if abs_max_identity >= similarity_bin >= abs_min_identity:
        # If there is no identity bin for the computed similarity_bin and it is higher than the max identity bin
        if similarity_bin not in identity_bins and similarity_bin >= max(identity_bins):
            # Assign the highest identity bin
            final_bin = max(identity_bins)
        # If there is still no identity bin for the computed similarity_bin, the similarity bin must be on the low end
        elif similarity_bin not in identity_bins:
            # Assign the lowest identity bin (which should be the abs_min_identity provided)
            final_bin = abs_min_identity
        # Otherwise the similarity_bin should match an available identity bin
        else:
            final_bin = similarity_bin
    return final_bin

def batch_iterator(iterator, batch_size):
    """
    Batch Iterator

    This can be used on any iterator, for example to batch SeqRecord objects from Bio.SeqIO.parse(...), or to batch
    Alignment objects from Bio.AlignIO.parse(...), or simply lines from a file handle. It is a generator function, and
    it returns lists of the entries from the supplied iterator.  Each list will have batch_size entries, although the
    final list may be shorter.

    Args:
        iterator (iterable): An iterable object, in this case a parser from the Bio package is the intended input.
        batch_size (int): The maximum number of entries to return for each batch (the final batch from the iterator may
        be smaller).
    Returns:
        list: A list of length batch_size (unless it is the last batch in the iterator in which case it may be fewer) of
        entries from the provided iterator.

    Scribed from Biopython (https://biopython.org/wiki/Split_large_file)
    """
    entry = True  # Make sure we loop once
    while entry:
        batch = []
        while len(batch) < batch_size:
            try:
                entry = next(iterator)
            except StopIteration:
                entry = None
            if entry is None:
                # End of file
                break
            batch.append(entry)
        if batch:
            yield batch

def filter_uniref_fasta(in_path, out_path):
    """
    Filter Uniref Fasta

    This function can be used to filter the fasta files provided by Uniref to remove the low quality sequences and
    sequences which represent protein fragments.

    Args:
        in_path (str): The path to the fasta file to filter (should be provided by Uniref so that the expected patterns
        can be found).
        out_path (str): The path to which the filtered fasta should be written.
    """
    sequences = []
    fragment_pattern = compile(r'^.*(\(Fragment\)).*$')
    low_quality_pattern = compile(r'^.*(LOW QUALITY).*$')
    record_iter = parse(open(in_path), "fasta")
    for i, batch in enumerate(batch_iterator(record_iter, 10000)):
        print('Batch: {}'.format(i))
        for seq_record in batch:
            fragment_check = fragment_pattern.search(seq_record.description)
            low_quality_check = low_quality_pattern.search(seq_record.description)
            if fragment_check or low_quality_check:
                continue
            sequences.append(seq_record)
            if len(sequences) == 1000:
                with open(out_path, 'a') as out_handle:
                    write(sequences=sequences, handle=out_handle, format='fasta')
                sequences = []

def characterize_alignment(file_name, query_id, abs_min_fraction=0.7, abs_max_identity=0.98, abs_min_identity=0.40,
                           interval=0.05):
    """
    Characterize Alignment

    This function seqs to characterize a given fasta formatted alignment in the same way which the DataSetGenerator uses
    when filtering sequences returned by the BLAST search.

    Args:
        file_name (str or path): The path to the file from which the alignment can be parsed. If a relative path is used
        (i.e. the ".." prefix), python's path library will be used to attempt to define the full path.
        query_id (str): The sequence identifier of interest.
        abs_min_fraction (float): The minimum fraction of the query sequence length for a passing hit.
        abs_max_identity (float): The absolute maximum identity for a passing hit.
        abs_min_identity (float): The absolute minimum identity for a passing hit.
        interval (int): The interval on which to define bins between min_identity and abs_min_identity in case
        sufficient sequences are not found at min_identity.
    Returns:
         float: The lowest fraction sequence length (expressed as a decimal) as compared to the query sequence.
         float: The highest fraction sequence length (expressed as a decimal) as compared to the query sequence.
         set: All fraction lengths observed in the provided alignment.
         dict: A dictionary mapping the identity bins, created by the provided abs_max_identity, abs_min_identity,
         min_identity, and interval values, to lists of sequence identifiers falling in that bin.
         dict: A dictionary mapping identity values, falling outside of the computed bins, to sequence identifiers
         having those identities when compared to the query sequence.
    """
    aln = SeqAlignment(file_name=file_name, query_id=query_id)
    aln.import_alignment()
    query_pos = aln.seq_order.index(query_id)
    full_query = str(aln.alignment[query_pos].seq)
    aligned_length = len(full_query)
    ungapped_length = len(str(aln.query_sequence).replace('-', ''))
    seq_fractions = {'Low': 0, 'Passing': 0, 'High': 0}
    max_fraction = 0.0
    min_fraction = 1.0
    identity_bins = [float(x / 100) for x in
                     range(int(100 * abs_min_identity), int(100 * abs_max_identity) + 1, int(100 * interval))]
    print(identity_bins)
    sequences = {x: [] for x in identity_bins}
    out_of_range = {}
    for i in range(aln.size):
        if aln.seq_order[i] == query_id:
            continue
        full_seq = str(aln.alignment[i].seq)
        ungapped_seq = len(full_seq.replace('-', ''))
        frac1 = ungapped_seq / float(ungapped_length)
        if frac1 < min_fraction:
            min_fraction = frac1
        if frac1 > max_fraction:
            max_fraction = frac1
        subject_fraction = min(frac1, ungapped_length / float(ungapped_seq))
        if subject_fraction < abs_min_fraction:
            if ungapped_seq < ungapped_length:
                seq_fractions['Low'] += 1
            else:
                seq_fractions['High'] += 1
            continue
        else:
            seq_fractions['Passing'] += 1
        id_count = 0
        for j in range(aligned_length):
            if full_query[j] == full_seq[j]:
                id_count += 1
        similarity_bin = determine_identity_bin(identity_count=id_count, length=aligned_length, interval=interval,
                                                abs_max_identity=abs_max_identity, abs_min_identity=abs_min_identity,
                                                identity_bins=set(identity_bins))
        if similarity_bin is None:
            identity = id_count / float(aligned_length)
            if identity not in out_of_range:
                out_of_range[identity] = []
            out_of_range[identity].append(aln.seq_order[i])
        else:
            sequences[similarity_bin].append(aln.seq_order[i])
    return min_fraction, max_fraction, seq_fractions, sequences, out_of_range

def parse_arguments():
    """
    Parse Arguments

    This method provides a nice interface for parsing command line arguments and includes help functionality.

    Returns:
        dict. A dictionary containing the arguments parsed from the command line and their arguments.
    """
    # Create input parser
    parser = argparse.ArgumentParser(description='Process DataSetGenerator command line arguments.')
    # Arguments for creating BLAST database input from a uniref fasta
    parser.add_argument('--custom_uniref', default=False, action='store_true',
                        help='Set this argument  to create a new fasta file as input to the makeblastdb tool.')
    parser.add_argument('--original_uniref_fasta', type=str, nargs='?',
                        help='The fasta to be filtered for the makeblastdb tool')
    parser.add_argument('--filtered_uniref_fasta', type=str, nargs='?',
                        help='The path where the filtered uniref fasta file should be saved.')
    # Arguments for characterizing an existing alignment
    parser.add_argument('--characterize_alignment', default=False, action='store_true',
                        help='Set this argument  to generate statistics on an alignment.')
    parser.add_argument('--file_name', type=str, nargs='?',
                        help="The path to the file from which the alignment can be parsed. If a relative path is used "
                             "(i.e. the '..' prefix), python's path library will be used to attempt to define the full "
                             "path.")
    parser.add_argument('--query_id', type=str, nargs='?',
                        help='The sequence identifier of interest.')
    # Arguments for creating a data set using the DataSet generator
    parser.add_argument('--create_data_set', default=False, action='store_true',
                        help='Set this argument to create a data set, i.e. download and generate all input files.')
    parser.add_argument('--input_dir', type=str, nargs='?', help='The path to the directory where the data for this '
                                                                 'data set should be stored. It is expected to contain '
                                                                 'at least one directory name ProteinLists which should'
                                                                 ' contain files where each line denotes a separate PDB'
                                                                 ' id (four letter code only).')
    parser.add_argument('--protein_list_fn', type=str, nargs='?',
                        help='The name of the file where the list of PDB ids (four letter codes) of which the data set '
                             'consists can be found. Each id is expected to be on its own line.')
    parser.add_argument('--processes', type=int, default=1, nargs=1,
                        help='The number of processes to use when performing the BLAST search.')
    parser.add_argument('--max_target_seqs', type=int, default=20000, nargs=1,
                        help='The maximum number of hits to look for in the BLAST database.')
    parser.add_argument('--e_value_threshold', type=float, default=0.05, nargs=1,
                        help='The maximum e-value for a passing hit.')
    parser.add_argument('--min_fraction', type=float, default=0.7, nargs=1,
                        help='The minimum fraction of the query sequence length for a passing hit.')
    parser.add_argument('--min_identity', type=float, default=0.40, nargs=1,
                        help='The absolute minimum identity for a passing hit.')
    parser.add_argument('--max_identity', type=float, default=0.98, nargs=1,
                        help='The absolute maximum identity for a passing hit.')
    parser.add_argument('--msf', type=bool, default=True, nargs=1,
                        help='Whether or not to create an msf version of the MUSCLE alignment.')
    parser.add_argument('--fasta', type=bool, default=True, nargs=1,
                        help='Whether or not to create an fasta version of the MUSCLE alignment.')
    parser.add_argument('--verbose', default=False, action='store_true',
                        help='Whether or not to print out information during data set generation')
    # Clean command line input
    arguments = parser.parse_args()
    arguments = vars(arguments)
    processor_count = cpu_count()
    arguments['processes'] = (arguments['processes'][0] if isinstance(arguments['processes'], list)
                              else arguments['processes'])
    arguments['max_target_seqs'] = (arguments['max_target_seqs'][0] if isinstance(arguments['max_target_seqs'], list)
                                    else arguments['max_target_seqs'])
    arguments['e_value_threshold'] = (arguments['e_value_threshold'][0]
                                      if isinstance(arguments['e_value_threshold'], list)
                                      else arguments['e_value_threshold'])
    arguments['min_fraction'] = (arguments['min_fraction'][0] if isinstance(arguments['min_fraction'], list)
                                 else arguments['min_fraction'])
    arguments['min_identity'] = (arguments['min_identity'][0] if isinstance(arguments['min_identity'], list)
                                 else arguments['min_identity'])
    arguments['max_identity'] = (arguments['max_identity'][0] if isinstance(arguments['max_identity'], list)
                                 else arguments['max_identity'])
    if arguments['processes'] > processor_count:
        arguments['processes'] = processor_count
    if arguments['custom_uniref']:
        if (not 'original_uniref_fasta' in arguments) or (not 'filtered_uniref_fasta'):
            raise ValueError('When custom_uniref is selected original_uniref_fasta and filtered_uniref_fasta must be '
                             'specfied.')
    if arguments['characterize_alignment']:
        if (not 'file_name' in arguments) or (not 'query_id' in arguments):
            raise ValueError('When characterize_alignment is selected file_name and query_id must be specified.')
    if arguments['create_data_set']:
        if (not 'input_dir' in arguments) or (not 'protein_list_fn' in arguments):
            raise ValueError('When create_data_set is selected input_dir and protein_list_fn must be specified.')
    return arguments

if __name__ == "__main__":
    args = parse_arguments()
    if args['custom_uniref']:
        filter_uniref_fasta(in_path=args['original_uniref_fasta'], out_path=args['filtered_uniref_fasta'])
    if args['characterize_alignment']:
        results = characterize_alignment(file_name=args['file_name'], query_id=args['query_id'],
                                         abs_min_fraction=args['min_fraction'],
                                         abs_max_identity=args['max_identity'],
                                         abs_min_identity=args['min_identity'])
        print('Sequence Fraction:\n\tMinimum:\t{}\n\tMaximum:\t{}'.format(results[0], results[1]))
        print('\t{} Sequence Fraction Low\n\t{} Sequence Fraction High\n\t{} Sequences Passed'.format(
            results[2]['Low'], results[2]['High'], results[2]['Passing']))
        print('\nSequence Identities:\n')
        for id_bin in sorted(results[3].keys()):
            print('\tBin_{}:\t{} Sequences'.format(id_bin, len(results[3][id_bin])))
        print('\nSequence Identities Outside Expected Range:\n')
        for id_bin in sorted(results[4].keys()):
            print('\tBin_{}:\t{} Sequences'.format(id_bin, len(results[4][id_bin])))
    if args['create_data_set']:
        generator = DataSetGenerator(input_path=args['input_dir'])
        generator.build_pdb_alignment_dataset(
            protein_list_fn=args['protein_list_fn'], processes=args['processes'],
            max_target_seqs=args['max_target_seqs'], e_value_threshold=args['e_value_threshold'],
            min_fraction=args['min_fraction'], min_identity=args['min_identity'],
            max_identity=args['max_identity'], msf=args['msf'], fasta=args['fasta'], verbose=args['verbose'])


import Stack
import consts
import pygame
import math
import BubblesGrid

screen = pygame.display.set_mode(
        (consts.WINDOW_WIDTH, consts.WINDOW_HEIGHT))

def draw_bubble(bubble):
    screen.blit(bubble["color"],(bubble["center_x"]-bubble["radius"],bubble["center_y"]-bubble["radius"]))

def draw_bubbles_popping(bubbles_popping):
    for bubble in bubbles_popping:
        draw_bubble(bubble)

def calc_mouse_angle(mouse_pos):
    x_diff = mouse_pos[0] - consts.ARROW_MIDBOTTOM_X
    y_diff = consts.ARROW_MIDBOTTOM_Y - mouse_pos[1]
    angle = math.degrees(math.atan2(y_diff, x_diff))
    return angle

def create_arrow(arrow_img):
    arrow = pygame.image.load(arrow_img)
    sized_arrow = pygame.transform.scale(arrow, (
        consts.ARROW_WIDTH, consts.ARROW_HEIGHT))

    # Create a box to put the arrow in, so that the rotation will be around
    # it's bottom (the box's center)
    arrow_box = pygame.Surface(
            (consts.ARROW_WIDTH, consts.ARROW_HEIGHT * 2), )
    arrow_box.fill(consts.BACKGROUND_COLOR)
    arrow_box.blit(sized_arrow, (0, 0))

    return arrow_box

def draw_arrow(arrow):
    rotated_arrow_rect = arrow.get_rect(
            center=(consts.ARROW_MIDBOTTOM_X, consts.ARROW_MIDBOTTOM_Y))
    screen.blit(arrow, rotated_arrow_rect)

frames = [pygame.image.load("lion-0001.png"),pygame.image.load("lion-0002.png"),pygame.image.load("lion-0003.png")]
frame = 0
counter = 0
def draw_shoot_animation():
    global  frame
    if counter % 10 == 0:
        frame+=1
    if frame > 2: frame = 0

    arrow = frames[frame]
    sized_arrow = pygame.transform.scale(arrow, (
        consts.ARROW_WIDTH, consts.ARROW_HEIGHT))

    # Create a box to put the arrow in, so that the rotation will be around
    # it's bottom (the box's center)
    arrow_box = pygame.Surface(
            (consts.ARROW_WIDTH, consts.ARROW_HEIGHT * 2), )

    draw_arrow(arrow_box)

def draw_border():
    line_y = consts.NUM_OF_LINES_LOSE * consts.BUBBLE_RADIUS * 2 - (
            consts.NUM_OF_LINES_LOSE - 1) * consts.ROWS_OVERLAP
    pygame.draw.line(screen, consts.BORDER_COLOR, start_pos=(0, line_y),
                     end_pos=(consts.WINDOW_WIDTH, line_y))

##########MINE#############
def draw_score():
    message = "SCORE: " + str(consts.score)
    draw_message(message, consts.TURNS_FONT_SIZE, consts.TURNS_COLOR, (10,consts.TURNS_LOCATION[1]-50))

    message = "HIGHEST SCORE: " + str(consts.highest_score)
    draw_message(message, consts.TURNS_FONT_SIZE, consts.TURNS_COLOR, (10, consts.TURNS_LOCATION[1] - 80))

def draw_turns(num_of_turns):
    message = consts.TURNS_TEXT + str(num_of_turns)
    draw_message(message, consts.TURNS_FONT_SIZE, consts.TURNS_COLOR,
                 consts.TURNS_LOCATION)

def draw_lose_message():
    draw_message(consts.LOSE_MESSAGE, consts.LOSE_FONT_SIZE,
                 consts.LOSE_COLOR, consts.LOSE_LOCATION)

def draw_win_message():
    draw_message(consts.WIN_MESSAGE, consts.WIN_FONT_SIZE,
                 consts.WIN_COLOR, consts.WIN_LOCATION)

def draw_message(message, font_size, color, location):
    font = pygame.font.SysFont(consts.FONT_NAME, font_size)
    text_img = font.render(message, True, color)
    screen.blit(text_img, location)

def draw_game(game_state):

    screen.fill(consts.BACKGROUND_COLOR)
    draw_arrow(game_state["rotated_arrow"])

    if game_state["is_bubble_fired"]:
        draw_bubble(game_state["bullet_bubble"])

    BubblesGrid.draw()
    draw_border()
    draw_turns(game_state["turns_left_to_add_row"])
    ##MINE##
    draw_score()
    Stack.draw()

    if len(game_state["bubbles_popping"]):
        BubblesGrid.animate_bubbles_pop(game_state["bubbles_popping"])
        draw_bubbles_popping(game_state["bubbles_popping"])

    elif game_state["state"] == consts.LOSE_STATE:
        draw_lose_message()
        pygame.display.flip()
        pygame.time.wait(1000*3)#3 seconds
        game_state["state"] = consts.RUNNING_STATE

    elif game_state["state"] == consts.WIN_STATE:
        draw_win_message()
        pygame.display.flip()
        pygame.time.wait(1000 * 3)  # 3 seconds
        game_state["state"] = consts.RUNNING_STATE

    pygame.display.flip()

################SRYYYYYYYYYYY####################
#################################################

pygame.init()
font = pygame.font.SysFont('Arial', 40)
class Button():
    def __init__(self, x, y, width, height, buttonText='Button', onclickFunction=None, onePress=False):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.onclickFunction = onclickFunction
        self.onePress = onePress
        self.alreadyPressed = False

        self.fillColors = {
            'normal': '#ffffff',
            'hover': '#666666',
            'pressed': '#333333',
        }

        self.buttonSurface = pygame.Surface((self.width, self.height))
        self.buttonRect = pygame.Rect(self.x, self.y, self.width, self.height)

        self.buttonSurf = font.render(buttonText, True, (20, 20, 20))
    def process(self):
        mousePos = pygame.mouse.get_pos()
        self.buttonSurface.fill(self.fillColors['normal'])
        if self.buttonRect.collidepoint(mousePos):
            self.buttonSurface.fill(self.fillColors['hover'])
            if pygame.mouse.get_pressed(num_buttons=3)[0]:
                self.buttonSurface.fill(self.fillColors['pressed'])
                if self.onePress:
                    self.onclickFunction()
                elif not self.alreadyPressed:
                    self.onclickFunction()
                    self.alreadyPressed = True
            else:
                self.alreadyPressed = False
        self.buttonSurface.blit(self.buttonSurf, [
            self.buttonRect.width / 2 - self.buttonSurf.get_rect().width / 2,
            self.buttonRect.height / 2 - self.buttonSurf.get_rect().height / 2
        ])
        screen.blit(self.buttonSurface, self.buttonRect)

RED = pygame.image.load("balls/red.png")
GREEN = pygame.image.load("balls/green.png")
BLUE = pygame.image.load("balls/blue.png")
ORANGE = pygame.image.load("balls/orange.png")
VIOLET = pygame.image.load("balls/violet.png")

run = True
chosen_list = []
easy_list = [pygame.image.load("balls/r.png"),pygame.image.load("balls/g.png"),pygame.image.load("balls/b.png")]
def easy():
    print("easy")
    global run
    global chosen_list
    consts.bubble_colors = easy_list.copy()
    chosen_list = easy_list
    run = False
med_list = [pygame.image.load("balls/violet2.png"), pygame.image.load("balls/green2.png"), pygame.image.load("balls/blue2.png"), pygame.image.load("balls/orange2.png")]
def med():
    global run
    consts.bubble_colors = med_list.copy()
    run = False
    global chosen_list
    chosen_list = med_list
hard_list = [pygame.image.load("balls/1.png"), pygame.image.load("balls/2.png"), pygame.image.load("balls/3.png"), pygame.image.load("balls/4.png"), pygame.image.load("balls/5.png")]
def hard():
    global run
    consts.bubble_colors = hard_list.copy()
    global chosen_list
    chosen_list = hard_list
    run = False

btns = []
btns.append(Button(30, 455, 200, 50, 'EASY', easy))
btns.append(Button(260, 455, 200, 50, 'MEDIUM', med))
btns.append(Button(490, 455, 200, 50, 'HARD', hard))
def Start_Screen()->bool:
    global run
    while run:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
                return False

        bg = pygame.image.load("balls/background4.png")
        screen.blit(bg,(0,0))

        for btn in btns:
            btn.process()

        pygame.display.flip()
    return True


from telegram import Bot, Update
from telegram.ext import Updater, Filters, CommandHandler, MessageHandler

from bot.config import TG_TOKEN, TG_API_URL

def start_bot(bot: Bot, update: Update):
	bot.send_message(
		chat_id=update.message.chat_id,
		text="Hi! Send me smth you want!"
	)

def do_echo(bot: Bot, update: Update):
	text = update.message.text
	bot.send_message(
		chat_id=update.message.chat_id,
		text=text,
	)

def main():
	bot = Bot(
		token=TG_TOKEN,base_url=TG_API_URL
	)
	updater = Updater(
		bot=bot
	)

	start_handler = CommandHandler("start", start_bot)
	message_handler = MessageHandler(Filters.text, do_echo)
	updater.dispatcher.add_handler(start_handler)
	updater.dispatcher.add_handler(message_handler)

	updater.start_polling()
	updater.idle()

if __name__ == '__main__':
	main()


from pathlib import Path

import typer

from unisport_abrechnung.bill import Bill
from unisport_abrechnung.configuration import Configuration
from unisport_abrechnung.template import Template
from unisport_abrechnung.util.date import parse_month
from unisport_abrechnung.util.file import find_free_file_name

def _unisport_abrechnung(base: Path, configuration: Configuration,
                         year: int, month: int, participant_counts: list[int]):
    # Create the bill data
    bill = Bill(configuration=configuration, year=year, month=month, participant_counts=participant_counts)

    # Determine file paths
    input_file = base / configuration.template.file
    output_file = find_free_file_name(base / f"{bill.default_file_name_stem()}.pdf")

    # Fill the template
    print(f"Reading template from {input_file}")
    writer = Template(input_file).fill(bill)

    # Write the output
    print(f"Writing output to {output_file}")
    assert not output_file.exists()
    with open(output_file, "wb") as f:
        writer.write(f)

def unisport_abrechnung(period:             str       = typer.Argument(""),
                        participant_counts: list[int] = typer.Argument(None)):

    # Load the configuration
    configuration_file = Path("unisport-abrechnung.toml")
    print(f"Loading configuration from {configuration_file}")
    configuration = Configuration.load(configuration_file)

    # If the period is not specified, query the user
    if not period:
        period = input("Billing period (mm/yyyy): ")

    # Parse the period
    year, month = parse_month(period)

    # If the participant counts are not specified, query the user (we need the
    # period for this)
    if not participant_counts:
        for day in configuration.class_.days(year, month):
            participant_count = int(input(f"Participant count for {year}-{month}-{day}: ") or "0")
            participant_counts.append(participant_count)

    # Run with cleaned up parameters
    _unisport_abrechnung(configuration_file.parent, configuration, year, month, participant_counts)

def main():
    typer.run(unisport_abrechnung)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Nov 20 12:43:07 2020

@author: kingsley
"""

#%% imports

import uproot
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt

from pylorentz import Momentum4
from pylorentz import Position4
from lbn import LBN, LBNLayer

# stop tensorflow trying to overfill GPU memory
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
  except RuntimeError as e:
    print(e)
#%% Data loading

tree = uproot.open("ROOTfiles/MVAFILE_AllHiggs_tt_pseudo.root")["ntuple"]

momenta_features = [ "pi_E_1", "pi_px_1", "pi_py_1", "pi_pz_1", #leading charged pi 4-momentum
              "pi_E_2", "pi_px_2", "pi_py_2", "pi_pz_2", #subleading charged pi 4-momentum
              "pi0_E_1","pi0_px_1","pi0_py_1","pi0_pz_1", #leading neutral pi 4-momentum
              "pi0_E_2","pi0_px_2","pi0_py_2","pi0_pz_2"] #subleading neutral pi 4-momentum

target = [    "aco_angle_1"]  #acoplanarity angle

selectors = [ "tau_decay_mode_1","tau_decay_mode_2",
             "mva_dm_1","mva_dm_2"
             ]

neutrino_features = [  "gen_nu_p_1", "gen_nu_p_2",
                       "gen_nu_phi_1", "gen_nu_phi_2",
                       "gen_nu_eta_1", "gen_nu_eta_2"]

met_features = ["met", "metx", "mety"]#, "metCov00", "metCov01", "metCov10", "metCov11"]

sv_features = ["sv_x_1", "sv_y_1", "sv_z_1",
               "sv_x_2", "sv_y_2", "sv_z_2"]
"""
               "svcov00_1", "svcov01_1", "svcov02_1",
               "svcov10_1", "svcov11_1", "svcov12_1",
               "svcov20_1", "svcov21_1", "svcov22_1",
               "svcov00_2", "svcov01_2", "svcov02_2",
               "svcov10_2", "svcov11_2", "svcov12_2",
               "svcov20_2", "svcov21_2", "svcov22_2"]
"""

ip_features = ["ip_x_1", "ip_y_1", "ip_z_1", "ip_x_2", "ip_y_2", "ip_z_2"]
"""
               "ipcov00_1", "ipcov01_1", "ipcov02_1", "ipcov10_1", "ipcov11_1", "ipcov12_1",
               "ipcov20_1", "ipcov21_1", "ipcov22_1", "ipcov00_2", "ipcov01_2", "ipcov02_2",
               "ipcov10_2", "ipcov11_2", "ipcov12_2", "ipcov20_2", "ipcov21_2", "ipcov22_2"]
"""

save_features = ["gen_phitt", "wt_cp_sm", "wt_cp_mm", "wt_cp_ps"]

df = tree.pandas.df(momenta_features+target+selectors
                    +neutrino_features+met_features+sv_features+sv_features+ip_features
                    +save_features)

df = df[
      (df["tau_decay_mode_1"] == 1)
    & (df["tau_decay_mode_2"] == 1)
    & (df["mva_dm_1"] == 1)
    & (df["mva_dm_2"] == 1)
]

#%% Clean up data

feature_quantities = momenta_features+met_features+sv_features+ip_features
target_quantities = neutrino_features

#Filter out -9999 values for neutrino momenta
df = df[df["gen_nu_p_1"] > -4000]
df = df[df["gen_nu_p_2"] > -4000]
df = df.dropna(subset=feature_quantities+target_quantities)

#%% train/eval split:

trainFrac = 0.5
df_train, df_eval = np.split(df, [int(trainFrac*len(df))], axis=0)

#%% Features and targets

def normalise(x):
    mu = tf.math.reduce_mean(x, axis=0)
    std = tf.math.reduce_std(x, axis=0)
    return (x-mu)/std, mu, std

#add visible product, met, sv, ip features
x_train = tf.convert_to_tensor(df_train[feature_quantities], dtype=tf.float32)

#y = tf.convert_to_tensor(df[neutrino_features], dtype=tf.float32)
y_train = tf.convert_to_tensor(df_train[neutrino_features], dtype=tf.float32)

#Eval data
x_eval = tf.convert_to_tensor(df_eval[feature_quantities], dtype=tf.float32)
y_eval = tf.convert_to_tensor(df_eval[neutrino_features], dtype=tf.float32)

#normalise
#x_train, x_mu, x_std = normalise(x_train)
y_train, y_mu, y_std = normalise(y_train)
#x_eval, _, _ = normalise(x_eval)
y_eval, _, _ = normalise(y_eval)

#%% Building network

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(),

    tf.keras.layers.Dense(64, activation="relu"),#, kernel_regularizer=tf.keras.regularizers.L1L2(0.01, 0.1)),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(64, activation="relu"),#, kernel_regularizer=tf.keras.regularizers.L1L2(0.01, 0.1)),
    #tf.keras.layers.Dense(200, activation="relu"),
    #tf.keras.layers.Dense(200, activation="relu"),
    #tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(6)
])

opt = tf.keras.optimizers.Adam(0.001)
loss_fn = tf.keras.losses.MeanSquaredError()
model.compile(optimizer=opt,
              loss=loss_fn,
              metrics=['mae', 'mse'])

print("Model compiled.")

#%% Training model

history = model.fit(x_train, y_train, validation_split=0.2, epochs=5)

#%% plot traning

plt.figure()
plt.plot(history.history['mse'], label="Training MSE")
plt.plot(history.history['val_mse'], label="Validation MSE")
plt.title("Loss on Iteration")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()

#%% Histograms

res = np.array((model(x_eval) - y_eval)*y_std + y_mu, dtype=np.float64)

plt.figure()
plt.title("NN Predicted Minus True p")
plt.xlabel("Phi / Radians")
plt.ylabel("Frequency")
#plt.xlim(-100, 100)
plt.hist(res[:,0], bins = 100, alpha = 0.5, label="p_1 mean={:.2f}, std={:.2f}".format(np.mean(res[:,0]), np.std(res[:,0])))
plt.hist(res[:,1], bins = 100, alpha = 0.5, label="p_2 mean={:.2f}, std={:.2f}".format(np.mean(res[:,1]), np.std(res[:,1])))
plt.grid()
plt.legend(loc="upper right")
plt.show()

plt.figure()
plt.title("NN Predicted Minus True Phi")
plt.xlabel("Phi / Radians")
plt.ylabel("Frequency")
#plt.xlim(-100, 100)
plt.xlim(-5, 5)
plt.hist(res[:,2], bins = 100, alpha = 0.5, label="phi_1 mean={:.2f}, std={:.2f}".format(np.mean(res[:,2]), np.std(res[:,2])))
plt.hist(res[:,3], bins = 100, alpha = 0.5, label="phi_2 mean={:.2f}, std={:.2f}".format(np.mean(res[:,3]), np.std(res[:,3])))
plt.grid()
plt.legend(loc="upper right")
plt.show()

plt.figure()
plt.title("NN Predicted Minus True Eta")
plt.xlabel("Phi / Radians")
plt.ylabel("Frequency")
#plt.xlim(-100, 100)
plt.xlim(-5, 5)
plt.hist(res[:,4], bins = 100, alpha = 0.5, label="eta_1 mean={:.2f}, std={:.2f}".format(np.mean(res[:,4]), np.std(res[:,4])))
plt.hist(res[:,5], bins = 100, alpha = 0.5, label="eta_2 mean={:.2f}, std={:.2f}".format(np.mean(res[:,5]), np.std(res[:,5])))
plt.grid()
plt.legend(loc="upper right")
plt.show()

#%% Add/remove columns to dataframe

#TODO: fix sv_variables
del df_eval['sv_x_1'], df_eval['sv_y_1'], df_eval['sv_z_1']
del df_eval['sv_x_2'], df_eval['sv_y_2'], df_eval['sv_z_2']

res = np.array(model(x_eval)*y_std + y_mu, dtype=np.float64)
df_eval['reco_nu_p_1'] = res[:,0]
df_eval['reco_nu_p_2'] = res[:,1]
df_eval['reco_nu_phi_1'] = res[:,2]
df_eval['reco_nu_phi_2'] = res[:,3]
df_eval['reco_nu_eta_1'] = res[:,4]
df_eval['reco_nu_eta_2'] = res[:,5]

#%%  Write root file

treeBranches = {column : str(df_eval[column].dtypes) for column in df_eval}
branchDict = {column : np.array(df_eval[column]) for column in df_eval}
tree = uproot.newtree(treeBranches, title="ntuple", compression=uproot.ZLIB(3))

with uproot.recreate("MVAFILE_AllHiggs_tt_reco.root") as f:
    f["ntuple"] = tree
    f["ntuple"].extend(branchDict)