import time

import pandas as pd
import numpy as np


def run_model(rna, tissue, rnaName='n/a'):

    setGlobals()
    k=5
    iterations=1000
    tissueNum = tissueToNum[tissue]

    rna = rna.upper()
    rbpConc = np.loadtxt(datadir+"/rbp_nM.csv", dtype = object, delimiter=",", usecols = tissueNum)
    prob = probMatrix(k, rna, rbpConc)

    sites = trial(k, rna, prob)
    diff = np.empty(iterations-1) #difference between subsequent samples
    for i in range(1, iterations):
        new = trial(k, rna, prob) #sampling step
        diff[i-1] = np.sum(np.absolute((sites+new)/(i+1)-sites/(i))) #calculate difference
        sites += new
    sites = sites/iterations #average
    return sites



def setGlobals(): #sets up array of rbp names, indexed kmers, kd table 
    global rbps, kmers, kd, kmerToNum, tissues, datadir, tissueToNum
    datadir = '/orange/ewang/RBPDB/src/data/'
    data = pd.read_csv('%skd_absolute.csv'%(datadir), sep=',', header=0, index_col=0)
    kmers = np.asarray(data.index)
    rbps = np.asarray(data.columns)
    kd = data.to_numpy()
    kmerToNum = {k: v for v, k in enumerate(kmers)}
    tissues = np.asarray(pd.read_csv('%srbp_copies.csv'%(datadir), sep=',', header=0, index_col=0).columns)
    hepg2 = np.asarray(pd.read_csv('%shepg2_copies.csv'%(datadir), sep=',', header=0, index_col=0).columns)
    tissues = np.concatenate((tissues, hepg2))
    tissueToNum = {k:v for v,k in enumerate(tissues)}
#    print(tissueToNum)



def probMatrix(k, rna, rbpConc): #creates probability table for each rbp-kmer interaction
    prob = np.empty((len(kmers),len(rbps)))
    concMatrix = np.broadcast_to(rbpConc,(len(kmers),len(rbpConc))) #for same shape as prob matrix
    rbpInd = np.arange(len(rbps)) #indices for vectorization
    kmerInd = np.arange(len(kmers))
    prob = prob.T
    prob[rbpInd] = concMatrix.T[rbpInd] #vectorize columns for rbp concentrations
    prob = prob.T
    prob[kmerInd] /= kd[kmerInd] #divide by Kd
    prob = np.hstack((prob,np.ones((1024,1)))) #append 1 for unbound state
    prob /= prob.sum(axis=1, keepdims=True) #divide by sum for probabilities
    return prob



def trial(k, rna, prob): #conducts one trial
    sites = np.zeros((len(rbps)+1, len(rna)), dtype = int)
    pos = np.arange(len(rna)-k+1) #possible binding positions
    np.random.shuffle(pos) #shuffle for randomness
    bound = set() #set of bound positions to prevent overlaps
    for i in pos:
        if i in bound:
            continue
        kmer = kmerToNum[rna[i:i+k]] #get index of kmer
        rbp = np.random.choice(np.arange(len(rbps)+1), p = prob[kmer]) #picks rbp
        if rbp == len(rbps): #if unbound, no changes in sites matrix
            sites[rbp][i] = 1
            continue
        sites[rbp][i:i+k] = 1 #state bound positions
        bound.update(set(range(i+1,i+k)))
    for i in np.arange(len(rna)-k+1, len(rna)):
        if i not in bound:
            sites[len(rbps)][i] = 1
    return sites