# -*- coding: utf-8 -*-

# <nbformat>3.0</nbformat>

# <codecell>

%load_ext autoreload

%autoreload 2

# <codecell>

from pandas import DataFrame, Series, merge, read_csv, MultiIndex, Index, concat

from subprocess import check_call

from tempfile import NamedTemporaryFile as NTF

import os, os.path

import numpy as np

from scipy.stats import ttest_ind

from itertools import groupby,combinations

from operator import itemgetter

from Bio import Phylo

import networkx

from random import shuffle

import csv, shlex, shutil

os.chdir('/home/will/Dropbox/HIVseqs/')

sys.path.append('/home/will/HIVReportGen/AnalysisCode/')

from SeqProcessTools import read_pat_seq_data, load_training_seq_data, align_seq_data_frame

# <codecell>

import glob

pat_files = glob.glob('/home/will/HIVReportGen/Data/PatientFasta/*.fasta')

pat_seq = read_pat_seq_data(pat_files, '/home/will/HIVReportGen/Data/BlastDB/ConBseqs.txt')

training_files = glob.glob('/home/will/HIVReportGen/Data/TrainingSequences/*.fasta')

training_data = load_training_seq_data(training_files)

align_lanl = align_seq_data_frame(training_data, '/home/will/HIVReportGen/Data/BlastDB/ConBseqs.txt')

all_seqs = concat([pat_seq, align_lanl])

# <codecell>

def get_pairwise_distances(seq_series, tree_file = None, seq_file = None):

return dmat

def extract_region(seq_series, start, stop):

nseqs = seq_series.map(lambda x: x[start:stop])

return nseqs

# <codecell>

def check_distance_pvals(mat_data, trop_dict):

return raw_pval, np.mean(cor_pvals), np.mean(g1dist), np.mean(g2dist), np.std(g1dist), np.std(g2dist)

# <codecell>

pssm_data = read_csv('/home/will/HIVReportGen/Data/TrainingSequences/pssm_data.csv', index_col = [0,1])

def decide_tropism(inval):

return np.nan

tropism_data = pssm_data['score'].map(decide_tropism).dropna()

trop_dict = {}

for (pat, visit), val in zip(tropism_data.index, tropism_data.values):

trop_dict[pat+'-'+visit] = val

benj_selected = []

with open('BensTropismLabels.csv') as handle:

reader = csv.DictReader(handle)

for row in reader:

trop_dict['%s-%s' % (row['Patient ID'], row['Visit'])] = row['Prediction'] == 'TRUE'

benj_selected.append((row['Patient ID'], row['Visit']))

benj_selected_index = MultiIndex.from_tuples(benj_selected, names = ['Patient ID', 'Visit number'])

# <codecell>

pure_seqs = MultiIndex.from_tuples([

], names = ['Patient ID', 'Visit number'])

equal_pure_seqs = MultiIndex.from_tuples([

], names = ['Patient ID', 'Visit number'])

# <codecell>

def make_tree_figure(wanted_seqs, trop_dict, tree_file):

networkx.draw_networkx(new_net, pos, with_labels = False, node_color = colors)

# <codecell>

check_regions = [#('Tat-seq-align', 'The Acidic domain', 0, 20),

]

indexes = [#('All LANL Seqs', tropism_data.index),

]

# <codecell>

from itertools import product

results = []

for (ind_name, inds), (seq_col, name, start, stop) in product(indexes, check_regions):

wanted = extract_region(all_seqs.ix[inds][seq_col].dropna(), start, stop)

#print wanted.index

#print('Treeing')

prot_name = seq_col.split('-')[0]

treename = 'fixeddomaintrees/%s-%s-%s-%i.nwk' % (ind_name, prot_name, name, start)

treename = treename.replace(' ', '-')

mat_data = get_pairwise_distances(wanted, tree_file=treename)

#print('Testing')

raw_p, cor_p, r5mean, x4mean, r5std, x4std = check_distance_pvals(mat_data, trop_dict)

if seq_col.startswith('LTR'):

start = start-454

stop = stop-454

print ind_name, name, start, raw_p, cor_p

results.append((ind_name, prot_name, name,start+1, stop+1, cor_p, r5mean, x4mean, r5std, x4std))

#pat_wanted = extract_region(pat_seq.ix[tropism_data.index][seq_col].dropna(), start, stop)

#fname = '/home/will/BenSeqs/Trees/' + name.replace(' ', '-')

#plt.figure(figsize = (20,20))

#make_tree_figure(wanted, trop_dict, fname + '.tree')

#plt.title(name)

#plt.savefig( fname + '_pure_seqs.png')

# <codecell>

with open('fixed_domain_analysis_with_new_ltr.csv', 'w') as handle:

writer = csv.writer(handle)

fields = ['Sequence Set', 'Protein Name', 'Domain Name', 'Region Start', 'Region Stop', 'p-value', 'R5-Mean', 'X4-Mean', 'R5-std', 'X4-std']

writer.writerow(fields)

writer.writerows(results)

# <codecell>

#widths = [5,10,15,20,25,30,35,45,50]

#indexes = [('BenJ Selected', benj_selected_index),

# ('All LANL Seqs', tropism_data.index),

# ]

#('Benj Pure Seq', pure_seqs)

#['Sequence Set', 'Protein Name', width, 'midpoint', 'p-value']

#large_results = []

#prots = [('Vpr-seq-align', 'Vpr',range(96))]

#for (ind_name, inds), width, (seq_col, prot, positions) in product(indexes, widths, prots):

#

# for mid in positions:

# print ind_name, width, prot, mid

# start = max(int(mid-(width/2)),0)

# stop = min(int(mid+(width/2)),positions[-1])

# wanted = extract_region(all_seqs.ix[inds][seq_col].dropna(), start, stop)

# mat_data = get_pairwise_distances(wanted)

# raw_p, cor_p, r5mean, x4mean = check_distance_pvals(mat_data, trop_dict)

#

# large_results.append((ind_name, prot, width,start+1, stop+1, cor_p, r5mean, x4mean))

# <codecell>

import contextlib

from tempfile import mkdtemp

@contextlib.contextmanager

def tmp_directory(*args, **kwargs):

pass

# <codecell>

from StringIO import StringIO

from subprocess import check_output

from Bio.Seq import Seq

from Bio import SeqIO

from Bio.SeqRecord import SeqRecord

from Bio.Seq import IUPAC

def write_nexus_alignment(seq_series, handle):

handle.write(strdata)

def write_mrbayes_commands(handle, alignment_path, output_path):

handle.write(cmd % tdict)

def run_mrbayes(cmd_path):

cmd = '/home/will/mb ' + cmd_path

check_output(shlex.split(cmd))

def reformat_nexus(inhandle, outhandle, trop_dict):

outhandle.write('end;\n')

def run_bats(formated_nexus_path, nreps = 5000):

return list(csv.DictReader(handle, fieldnames=headers, delimiter = '\t'))[:-2]

def run_MrBats_analysis(seq_series, trop_dict, tree_file):

return out

# <codecell>

from copy import deepcopy

from concurrent.futures import ThreadPoolExecutor

from itertools import chain, imap, product

indexes = [#('BenJ Selected', benj_selected_index),

]

def linker_code(tup):

return final

#(ind_name, prot_name, name,start+1, stop+1, cor_p, r5mean, x4mean)

new_method_inputs = []

for (ind_name, inds), (seq_col, name, start, stop) in product(indexes, check_regions):

wanted = extract_region(all_seqs.ix[inds][seq_col].dropna(), start, stop)

#print wanted.index

#print('Treeing')

prot_name = seq_col.split('-')[0]

treename = 'newdomaintrees/%s-%s-%s-%i.nwk' % (ind_name, prot_name, name, start)

extra_dict = {

'IndName':ind_name,

'ProtName':prot_name,

'Domain':name,

'Start':start,

'Stop':stop,

}

treename = treename.replace(' ', '-')

#print treename, len(wanted)

if len(wanted)>10:

new_method_inputs.append((wanted.copy(), treename, deepcopy(extra_dict)))

#raise KeyError

#results_so_far = []

with ThreadPoolExecutor(max_workers = 30) as executor:

res = executor.map(linker_code, new_method_inputs)

for row in chain.from_iterable(res):

print row['Domain'], row['IndName'], row['Statistic'] ,row['significance']

#results_so_far.append(row)

# <codecell>

tmp = DataFrame(results_so_far)

tmp.to_csv('new_method_results.csv')

# <codecell>

from itertools import islice

widths = [5,10,15,20,25,30,35,45,50]

indexes = [('BenJ Selected', benj_selected_index),

('All LANL Seqs', tropism_data.index),

('Benj Pure Seq', pure_seqs),

]

prots = [('Vpr-seq-align', 'Vpr',range(96))]

def take(n, iterable):

"Return first n items of the iterable as a list"

return list(islice(iterable, n))

def yield_regions(indexes, widths, prots):

yield (wanted.copy(), None, deepcopy(extra_dict))

#['Sequence Set', 'Protein Name', 'width', 'midpoint', 'p-value']

vpr_window_results = []

block_size = 500

with ThreadPoolExecutor(max_workers = 30) as executor:

iterable = yield_regions(indexes, widths, prots)

block = take(block_size, iterable)

while block:

res = executor.map(linker_code, block)

for row in chain.from_iterable(res):

print row['width'], row['IndName'], row['Statistic'] ,row['significance']

vpr_window_results.append(row)

block = take(block_size, iterable)

# <codecell>

windowed_df = DataFrame(vpr_window_results)

# <codecell>

windowed_df.head()

# <codecell>

import dendropy

fixtrees = glob.glob('newdomaintrees/*.nwk')

for f in fixtrees:

if 'Equal' not in f:

continue

with open(f) as handle:

tree = dendropy.Tree.get_from_stream(open(f), 'nexus')

tree.deroot()

rmnodes = [tree.prune_subtree(t, update_splits = True) for t in tree.leaf_nodes() if t.get_node_str().endswith("copy'")]

#tree.prune_taxa(rmnodes)

nf = f.replace('newdomaintrees', 'unrootedtrees-equal')

with open(nf, 'w') as handle:

tree.write_to_stream(handle, 'newick')

# <codecell>

tmp = list(tree.leaf_nodes())

# <codecell>

labels = [n.get_node_str() for n in rmnodes]

tree.prune_taxa_with_labels(labels, update_splits = True)

# <codecell>

[tree.prune_subtree(t) for t in tree.leaf_nodes() if t.get_node_str().endswith("copy'")]

# <codecell>

with open('/home/will/tmpstuf/tmpyw6FXN/seqalign.nxs.trprobs') as handle:

treeL = dendropy.TreeList.get_from_stream(handle, 'nexus')

# <codecell>

print treeL[0].description()

# <codecell>

with open('/home/will/tmpstuf/tmpyw6FXN/tmptree.nxs', 'w') as handle:

treeL.write_to_stream(handle, 'nexus')

# <codecell>

for tree in treeL:

for leaf in tree.leaf_iter():

print str(leaf.taxon)

# <codecell>