# -*- coding: utf-8 -*-
"""
Estimate actual counts of barcode combos from observed counts 
based on combinations and probability of non-binding (p=0.494).
"""

# console command
# nohup python3 -u count_corrections.py 2 1 > 2.1.out &
# ps xw
# kill <pid>

import pandas as pd
from scipy.special import comb
import itertools
import sys
import os

# process one sample and subsample at a time; can concatenate later
sample_id = int(sys.argv[1]) # first argument
subsample_id = int(sys.argv[2]) # second argument

p = .494 # probability of adjacent binding barcodes
L_max = 5

num_barcode_types = 34

input_filename = 'Barcode combos'

df = pd.read_csv(input_filename+'.csv')
df = df[['Sample','Subsample','Node','# Mers','ObservedCount']] # subsamples are for bootstrap sampling
df = df.set_index(['Sample','Subsample','Node','# Mers'])
df = df.sort_index()

df['ActualCount'] = [0 for r in range(len(df))]

def complex_name2barcodes(complex_name):
    complex_barcodes = complex_name[1:].split('C') # remove first 'C' and split into each barcode
    return list(map(int, complex_barcodes))

def complex_barcodes2name(complex_barcodes):
    complex_name = 'C' + 'C'.join(str(b).zfill(2) for b in complex_barcodes)
    return complex_name

# regex to match if full seq has all barcodes in the subseq
def complex_barcodes2name_regex(complex_barcodes):
    complex_name = '.*C'.join(str(b).zfill(2) for b in complex_barcodes) # assumes combos are ordered ascending
    complex_name = '.*C' + complex_name + '.*'
    return complex_name

barcode_list = [i for i in range(1,num_barcode_types+1)]
barcode_name_list = [complex_barcodes2name([i]) for i in range(1,num_barcode_types+1)]

props = [[0 for i in range(L_max)] for j in range(L_max)] # initialize 2D array to store props with [l, L] indices


# proportion of actual barcode combos observed
def get_proportion(l, L):
    if props[l-1][L-1] > 0:
        prop = props[l-1][L-1]
    
    else:    
        alpha = l / comb(L, l-1)
        beta = 2 / comb(L-(l-1), 1)
        prop = alpha * pow(p, l-1) * pow(1-p, L-l>=1) * pow(beta+(1-beta)*(1-p), L-l>=2)
        props[l-1][L-1] = prop
        
    return prop

# count num of subseq from seq of length L
def get_subseq_measured_counts(subseq, L, sample, subsample):
    subseq_barcodes = complex_name2barcodes(subseq)
    l = len(subseq_barcodes) # num_subseq_barcodes
    
    count = 0
    
    if l > L:
        count = 0
    
    elif l == L:
        if (sample,subsample,subseq,L) in df.index:
            obs_count = df.loc[(sample,subsample,subseq,L),('ObservedCount')]
        else:
            obs_count = 0

        count = obs_count - sum([get_subseq_measured_counts(subseq, L, sample, subsample) for L in reversed(range(l+1,L_max+1))])
        seq = seq_only = subseq + ''
        
    elif (sample,subsample,subseq,L) in df.index: # already calculated, so just return it
       count = df.loc[(sample,subsample,subseq,L),('ObservedCount')]
       return count
        
    else: # l < L
        count = get_proportion(l, L) * get_actual_count(subseq, L, sample, subsample)    
        num_barcodes_any = L - l
        seq = seq_any = subseq + '+'.join('' for a in range(num_barcodes_any+1)) # append '+' for each extra mer
        
    df.loc[(sample,subsample,seq,L),('AttributedCount')] = (count)
    return count

# derive actual count of protein combos from measured counts of shorter barcode combos
def get_actual_count(seq, L, sample, subsample):
    subseq_barcodes = complex_name2barcodes(seq)
    l = len(subseq_barcodes) # num_subseq_barcodes
    
    actual_count = 0
            
    if l > L: 
        # protein combo length longer than barcode combo length 
        # => not possible to get barcode comboes from these protein combos
        return 0 # do nothing        
      
    elif l == L:
        obs_count = get_subseq_measured_counts(seq, L, sample, subsample)
        actual_count = obs_count / (get_proportion(l, L))
        seq = seq + ''
        
    elif (sample,subsample,seq,L) in df.index: # already calculated, so just return it
       count = df.loc[(sample,subsample,seq,L),('ActualCount')]
       return count
        
    else: # l < L
        # protein combo length shorter than barcode combo length 
        # should already have been calculated and cached
        
        # only handle subset with matching L
        # collect counts of protein combos that are derived from smaller observed barcode combos
        
        # replace index with seq and 1 extra char per barcode, then count if length increased by l to match
        
        match_sample = df.index.get_level_values('Sample') == sample
        match_subsample = df.index.get_level_values('Subsample') == subsample        
        
        subseq_barcodes_regex = complex_barcodes2name_regex(subseq_barcodes)
        match_node = (df.index.get_level_values('Node').str.contains(subseq_barcodes_regex) & ~df.index.get_level_values('Node').str.endswith('+'))
                
        match_Lmer = df.index.get_level_values('# Mers') == L
        _df = df[match_sample & match_subsample & match_node & match_Lmer]
        actual_count = _df['ActualCount'].sum()
        
        num_barcodes_any = L - l
        seq = seq + '+'.join('' for a in range(num_barcodes_any+1)) # append ('+' for each extra mer) and rename seq

    df.loc[(sample,subsample,seq,L),('ActualCount')] = (actual_count)
    return actual_count


"""
Iterate through samples and count barcode combos by L-mers
"""

if not os.path.isdir(input_filename):
    os.mkdir(input_filename)

for sample in [sample_id]:
    print('Processing sample =', sample)
    
    for subsample in [subsample_id]:
        print('Processing subsample =', subsample)
        
        # count longer Lmers first before shorter ones
        for L in reversed(range(1,5+1)):
            print('Processing L =',L)
            match_sample = df.index.get_level_values('Sample') == sample
            match_subsample = df.index.get_level_values('Subsample') == subsample
            match_Lmer = df.index.get_level_values('# Mers') == L
            _df = df[match_sample & match_subsample & match_Lmer]
            i = 0
            for (sample,subsample,seq,L) in list(_df.index.values): # iterating to get seq
                get_actual_count(seq, L, sample, subsample)
                i += 1
                if i % 100 == 0: # print at every few rows to track progress
                    print(i,seq)
            
            match_sample = df.index.get_level_values('Sample') == sample
            match_subsample = df.index.get_level_values('Subsample') == subsample
            __df = df[match_sample & match_subsample]
            __df.to_csv(input_filename+'/ActualCounts_S'+str(sample)+'.'+str(subsample)+'_'+str(L)+'mer_processed.csv', sep=',')
        
        for seq in barcode_name_list:
            get_actual_count(seq, 1, sample, subsample)

_df = df[~((df['AttributedCount'] == 0) & (df['ActualCount'] == 0))]
pd.options.display.float_format = '{:,.2f}'.format
print(_df)
_df.to_csv(input_filename+'/S'+str(sample_id)+'.'+str(subsample_id)+'_'+'actual_counts.csv', sep=',')