np.random.seed(123)
torch.manual_seed(123)

"""
Explanation: Class balance estimation from weak supervision
In this notebook, we'll use a special class, ClassBalanceModel, for estimating the class balance, i.e. $P(Y)$, from unlabeled data and a set of conditionally independent labeling functions.  We rely on this conditional independence assumption in this approach (note that these LFs could be a subset of the larger set of LFs used in the LabelModel, selected to satisfy this assumption).
Note that this is an advanced tutorial, and thus will go into slightly more detail in certain aspects, and skip over some basics in other parts.
Note also that:
* The LFs in this tutorial do not abstain, however the ClassBalanceModel does handle abstaining LFs.
* We directly generate a three-way overlaps matrix $O$ as a toy synthetic example; however, the ClassBalanceModel can also take in an empirical label matrix $L$ and form $O$ internally.
For further details, see the tests in tests/metal/label_model/test_class_balance.py.
End of explanation
"""


K = 3

# Generate the true class balance to be recovered
p_Y = np.random.random(K)
p_Y /= p_Y.sum()
p_Y

"""
Explanation: Step 1: Generate the true class balance to be recovered
End of explanation
"""


M = 25
Cs = []
for i in range(M):
    a = np.random.random((K, K)) + (K-1)*np.eye(K)
    Cs.append( a @ np.diag(1 / a.sum(axis=0)) )
C = np.array(Cs)
assert np.all(np.abs(C.sum(axis=1) - 1) < 1e-5)
C[0]

"""
Explanation: Step 2: Generate the true conditional probability tables (CPTs) for the LFs
We start by generating the conditional probability table $C$ for $M$ labeling functions (LFs) which we assume are conditionally independent given the latent true label $Y$, i.e. $\lambda_i \perp \lambda_{j \neq i} | Y$.
We have:
$$
C_{i,y',y} = P(\lambda_i = y' | Y = y)
$$
Note that this table should be normalized such that:
$$
\sum_{y'} C_{i,y',y} = 1
$$
Note also that we skew the diagonal entries to roughly enforce our other assumption here: that the LFs are (on average) better than random. We will later rely on this assumption for breaking a final column-wise symmetry in our class balance estimation problem.
End of explanation
"""


# Compute mask
mask = torch.ones((M, M, M, K, K, K)).byte()
for i, j, k in product(range(M), repeat=3):
    if len(set((i, j, k))) < 3:
        mask[i, j, k, :, :, :] = 0

"""
Explanation: Step 3: Generate the three-way overlaps tensor $O$ of conditionally-independent LFs
Now we can directly generate $O$.
By our conditional independence assumption, for distinct $i,j,k$, we have:
$$
P(\lambda_i = y', \lambda_j = y'', \lambda_k = y''' | Y = y) = C_{i,y',y} C_{j,y'',y} C_{k,y''',y}
$$
Thus we have:
$$
O_{i,j,k,y',y'',y'''} = \sum_y P(Y=y) C_{i,y',y} C_{j,y'',y} C_{k,y''',y}
$$
Note that we explicitly mask out entries for which $i,j,k$ are not distinct; however, note that this is redundant given that this masking also occurs in the ClassBalanceModel class we will use!
End of explanation
"""


%%time
O = np.einsum('aby,cdy,efy,y->acebdf', C, C, C, p_Y)
O = torch.from_numpy(O).float()
O[1-mask] = 0

"""
Explanation: Note that we use einsum, which is great!
End of explanation
"""


%%time
from metal.label_model.class_balance import ClassBalanceModel

model = ClassBalanceModel(K, abstains=False)
model.train_model(O=O, verbose=False)

print(f"Estimated class balance: {model.class_balance}")
print(f"True class balance: {p_Y}")
print(f"Avg. error: {np.mean(np.abs(model.class_balance - p_Y))}")

"""
Explanation: Step 4: Recover the class balance $P(Y)$ using ClassBalanceModel
End of explanation
"""


