# set up the environment
import os
import sys
import gzip
from functools import reduce
import tempfile

import pandas as pd

from Bio import SeqIO

import PhiSpyModules

"""
Explanation: PhiSpy
This is a Jupyter Notebook that shows how to run PhiSpy manually. You can run through all the steps that PhiSpy takes to determine whether a genome contains a prophage, and inspect all of the data generated by PhiSpy.
You will need to install PhiSpy, Jupyter Notebooks, and pandas
<small>Note: PhiSpy does not normally use pandas, but we use it here to visualize the data!</small>
End of explanation
"""


print("PhiSpy version: " + PhiSpyModules.__version__)

"""
Explanation: Check the PhiSpy version
We recommend at least version 4.0.3 but preferable 4.1 or higher
End of explanation
"""


genbankfile = "../test_genbank_files/Yersinia_pestis_KIM.gb.gz"

"""
Explanation: Define your genbank file here.
You may need to set the full path to the file. You can use gzip compressed file or an uncompressed file. Obviously, you will know whether it is compressed or not, but this demonstrates how PhiSpy determines when a file is compressed.
This should be the only line you need to change to run PhiSpy completely!
End of explanation
"""


min_contig_size = 1000

if PhiSpyModules.is_gzip_file(genbankfile):
    handle = gzip.open(genbankfile, 'rt')
else:
    handle = open(genbankfile, 'r')

record  = PhiSpyModules.SeqioFilter(filter(lambda x: len(x.seq) > min_contig_size, SeqIO.parse(handle, "genbank")))
handle.close()

# we check to make sure there are some contigs left to process
ncontigs = reduce(lambda sum, element: sum + 1, record, 0)
print(f"There are {ncontigs} contigs to predict prophages on!")

"""
Explanation: Parse the file
We use BioPython to parse the file, but we also add a few additional mehtods to the standard BioPython object to ease parsing. We also merge or split compound features (those with more than one location along the chromosome) to appropriately handle them.
Our record object is an extended SeqIO.parse object.
End of explanation
"""


parameters = {
    'kmers_type': "all",
    'window_size': 30,
    'record': record,
    'expand_slope': False,
    'training_set': "data/trainSet_genericAll.txt",
    'randomforest_trees': 5,
    'threads': 4,
    'quiet': True,
    'nonprophage_genegaps': 10,
    'number': 5,
    'color' : True,
    'evaluate':False,
    'make_training_data':None,
    'skip_search':False,
    'phmms':None,
    'phage_genes' : 2,
    'metrics': ['orf_length_med', 'shannon_slope', 'at_skew', 'gc_skew', 'max_direction']
}

"""
Explanation: Define the parameters that we will use
These are normally provided as command line options, but for jupyter we set them here
paramter | meaning | options | default value
--- | --- | --- | --
kmers_type | What do we count kmers with? | all, codon, simple | all
window_size | How many consecutive ORFs to include? | an integer | 30
record | the Bio.SeqIO object with all the sequences | a Bio.SeqIO object | record
expand_slope | whether to use the square of the slope of the Shannon scores | True or False | False
number | Number of consecutive genes in a region of window size that must be prophage genes | an integer | 5
nonprophage_genegaps | The number of non phage genes betweeen prophages | an integer | 10
quiet | Don't make additional outputs | True or False | True
Note: You can add an additional paramter, make_training_data here (its actual value doesn't matter) that will append an additional column to the output for each ORF that includes a 1 (True) if the ORF is thought or stated to be a phage gene or 0 (False) otherwise.
End of explanation
"""


parameters['test_data'] = PhiSpyModules.measure_features(**parameters)
# note that if you include make_training_data you will need to add an "is_phage" column here
test_df = pd.DataFrame(parameters['test_data'], columns = ['orf_length_med', 'shannon_slope', 'at_skew',
                                             'gc_skew', 'max_direction', 'phmms'])
test_df.head()

"""
Explanation: Generate the test data
This is the step that actually does all the measurements!
In this example, we convert the output to a pandas dataframe for visualization and exploration.
End of explanation
"""


parameters['rfdata'] = PhiSpyModules.call_randomforest(**parameters)
parameters['initial_tbl'] = PhiSpyModules.make_initial_tbl(**parameters)
parameters['output_dir'] = tempfile.mkdtemp()
initial_table_df = pd.DataFrame(parameters['initial_tbl'], columns = ['gene id', 'function', 'contig', 'start', 'stop', 'position', 'rank', 'my status', 'pp'])
initial_table_df.head()

"""
Explanation: Run the random forest
Here we run the random forest to identify the phages, and combine that into our initial table as the rank column.
End of explanation
"""


parameters['pp'] = PhiSpyModules.fixing_start_end(**parameters)
pp_df = pd.DataFrame.from_dict(parameters['pp']).transpose()
pp_df

"""
Explanation: Refine the predictions
Finally, we refine the predictions from the random forest and other metrics, and then predict the att sequences.
End of explanation
"""


parameters['final_tbl'] = []
for i in parameters['initial_tbl']:
    my_fs = PhiSpyModules.evaluation.check_pp(i[2], i[3], i[4], parameters['pp'])
    parameters['final_tbl'].append(i + [my_fs])

    final_df  = pd.DataFrame(parameters['final_tbl'], columns = ['gene id', 'function', 'contig', 'start', 'stop', 'position', 'rank', 'my status', 'pp', 'final status'])
final_df.head()

"""
Explanation: Our pp_df data frame has our final prophage predictions for this genome! 
Make the final table
Here we just append the pp number of the prophage to the table to show what are pp regions in the data frame.
End of explanation
"""


