# Module for using BLOSUM62 matrix 
import numpy as np 
np.set_printoptions(suppress=True)

##############################################################
#### Create BLOSUM62 s_ij matrix [=] bits (NOT half bits) ####
##############################################################
# A        R        N        D        C        Q        E        G        H        I        L        K        M        F        P        S        T        W        Y        V
b62_sij_str = """
  1.9646
 -0.7068   2.7367
 -0.7654  -0.2199   2.8266
 -0.8767  -0.8029   0.6358   2.8871
 -0.2043  -1.6946  -1.3299  -1.7300   4.2911
 -0.4020   0.4914   0.0008  -0.1567  -1.4509   2.6426
 -0.4319  -0.0577  -0.1340   0.7552  -1.8062   0.9273   2.4514
  0.0798  -1.1521  -0.2114  -0.6568  -1.2502  -0.8926  -1.0551   2.7816
 -0.8126  -0.1249   0.2892  -0.5595  -1.4939   0.2240  -0.0588  -1.0204   3.7555
 -0.6609  -1.4951  -1.6085  -1.5606  -0.6138  -1.3848  -1.5972  -1.8624  -1.6158   1.9993
 -0.7323  -1.0773  -1.6895  -1.8028  -0.6387  -1.0670  -1.4232  -1.8135  -1.3934   0.7608   1.9247
 -0.3670   1.0544  -0.0895  -0.3509  -1.5182   0.6363   0.3877  -0.7640  -0.3605  -1.3351  -1.2234   2.2523
 -0.4676  -0.6836  -1.0754  -1.5293  -0.7099  -0.2105  -0.9990  -1.3383  -0.7756   0.5634   0.9959  -0.6774   2.6963
 -1.1050  -1.3932  -1.4970  -1.7419  -1.1877  -1.5822  -1.5962  -1.5537  -0.6171  -0.0804   0.2074  -1.5393   0.0063   3.0230
 -0.4071  -1.0543  -1.0002  -0.7401  -1.3976  -0.6410  -0.5581  -1.0668  -1.0805  -1.3783  -1.4300  -0.5068  -1.2382  -1.7986   3.6823
  0.5579  -0.3824   0.3005  -0.1305  -0.4375  -0.0506  -0.0735  -0.1462  -0.4408  -1.1741  -1.2213  -0.1017  -0.7404  -1.1845  -0.4045   1.9422
 -0.0227  -0.5612  -0.0230  -0.5254  -0.4333  -0.3377  -0.4316  -0.7877  -0.8429  -0.3588  -0.5987  -0.3348  -0.3331  -1.0538  -0.5376   0.6906   2.2727
 -1.2634  -1.3397  -1.8480  -2.1072  -1.1521  -0.9732  -1.4177  -1.2457  -1.1711  -1.2903  -0.8159  -1.4782  -0.7124   0.4588  -1.8271  -1.3759  -1.2145   5.2520
 -0.8820  -0.8469  -1.0409  -1.5325  -1.2036  -0.7105  -1.0102  -1.5199   0.8463  -0.6657  -0.5310  -0.9100  -0.4974   1.4696  -1.4599  -0.8429  -0.8030   1.0771   3.2975
 -0.0947  -1.2513  -1.4382  -1.5713  -0.4038  -1.0992  -1.2211  -1.5694  -1.5587   1.2735   0.3942  -1.1312   0.3436  -0.4245  -1.1744  -0.8231  -0.0278  -1.4171  -0.6038   1.8845
 """
b62_sij = np.full((20,20), np.nan)
b62_sij_rows = [[float(a) for a in row.split()] for row in b62_sij_str.split('\n')[1:-1]]
for i,r in enumerate(b62_sij_rows):
    for j,sij in enumerate(r):
        b62_sij[i,j] = sij
        b62_sij[j,i] = sij # symmetric 



###############################################
#### Marginal Probabilities by AA = p(i,*) ####
###############################################
# A     R     N     D     C     Q     E     G     H     I     L     K     M     F     P     S     T     W     Y     V
marginals = "0.074 0.052 0.045 0.054 0.025 0.034 0.054 0.074 0.026 0.068 0.099 0.058 0.025 0.047 0.039 0.057 0.051 0.013 0.032 0.073" 
marginals = np.array([float(a) for a in marginals.split()])



#NOTE: What is the difference between marginals and freq_i? not sure 
################################
#### AA Probabilities = fi: ####
################################ 
# A     R     N     D     C     Q     E     G     H     I     L     K     M     F     P     S     T     W     Y     V
freq_i="0.078 0.051 0.041 0.052 0.024 0.034 0.059 0.083 0.025 0.062 0.092 0.056 0.024 0.044 0.043 0.059 0.055 0.014 0.034 0.072"
freq_i=np.array([float(a) for a in freq_i.split()])



##########################
#### Joint probs q_ij ####
##########################
# A      R      N      D      C      Q      E      G      H      I      L      K      M      F      P      S      T      W      Y      V
b62_qij_str = """
0.0215
0.0023 0.0178
0.0019 0.0020 0.0141
0.0022 0.0016 0.0037 0.0213
0.0016 0.0004 0.0004 0.0004 0.0119
0.0019 0.0025 0.0015 0.0016 0.0003 0.0073
0.0030 0.0027 0.0022 0.0049 0.0004 0.0035 0.0161
0.0058 0.0017 0.0029 0.0025 0.0008 0.0014 0.0019 0.0378
0.0011 0.0012 0.0014 0.0010 0.0002 0.0010 0.0014 0.0010 0.0093
0.0032 0.0012 0.0010 0.0012 0.0011 0.0009 0.0012 0.0014 0.0006 0.0184
0.0044 0.0024 0.0014 0.0015 0.0016 0.0016 0.0020 0.0021 0.0010 0.0114 0.0371
0.0033 0.0062 0.0024 0.0024 0.0005 0.0031 0.0041 0.0025 0.0012 0.0016 0.0025 0.0161
0.0013 0.0008 0.0005 0.0005 0.0004 0.0007 0.0007 0.0007 0.0004 0.0025 0.0049 0.0009 0.0040
0.0016 0.0009 0.0008 0.0008 0.0005 0.0005 0.0009 0.0012 0.0008 0.0030 0.0054 0.0009 0.0012 0.0183
0.0022 0.0010 0.0009 0.0012 0.0004 0.0008 0.0014 0.0014 0.0005 0.0010 0.0014 0.0016 0.0004 0.0005 0.0191
0.0063 0.0023 0.0031 0.0028 0.0010 0.0019 0.0030 0.0038 0.0011 0.0017 0.0024 0.0031 0.0009 0.0012 0.0017 0.0126
0.0037 0.0018 0.0022 0.0019 0.0009 0.0014 0.0020 0.0022 0.0007 0.0027 0.0033 0.0023 0.0010 0.0012 0.0014 0.0047 0.0125
0.0004 0.0003 0.0002 0.0002 0.0001 0.0002 0.0003 0.0004 0.0002 0.0004 0.0007 0.0003 0.0002 0.0008 0.0001 0.0003 0.0003 0.0065
0.0013 0.0009 0.0007 0.0006 0.0003 0.0007 0.0009 0.0008 0.0015 0.0014 0.0022 0.0010 0.0006 0.0042 0.0005 0.0010 0.0009 0.0009 0.0102
0.0051 0.0016 0.0012 0.0013 0.0014 0.0012 0.0017 0.0018 0.0006 0.0120 0.0095 0.0019 0.0023 0.0026 0.0012 0.0024 0.0036 0.0004 0.0015 0.0196
"""
b62_qij_rows = [[float(a) for a in row.split()] for row in b62_qij_str.split('\n')[1:-1]]
b62_qij = np.full((20,20), np.nan)
for i,r in enumerate(b62_qij_rows):
    for j,qij in enumerate(r):
        b62_qij[i,j] = qij
        b62_qij[j,i] = qij # symmetric


########################################
#### Conditional probability p(A|B) ####
########################################
"""
To calculate p(A|B), do this:
    1. take 2^(b62_sij) element wise to get matrix of odds ratios 
    2. multiply result by marginals (row wise)
    3. the result is p(col|row)
"""
# probability of j (col) conditioned on i (row)
p_j_given_i = 2**b62_sij.astype(float)
p_j_given_i *= marginals[None,:]          # casts along the row 
p_j_given_i /= p_j_given_i.sum(axis=-1, keepdims=True)