#! /usr/bin/env python
import numpy as np
import pytorch_lightning as pl
import pytorch_lightning.metrics
import torch

# To demonstrate the bugs in the current implementation of the
# ConfusionMatrix computation, I am going to use a contrived 4 class
# example.  The data will consist of 20 samples - 9 of class 0, 0 of class
# 1, 10 of class 2, and 1 of class 3.  The data is arranged such that the
# true class is the first column and the declared class is the second
# column.
data = np.array([
    [3, 3],
    [0, 0],
    [0, 0],
    [0, 0],
    [0, 0],
    [0, 1],
    [0, 2],
    [0, 2],
    [0, 2],
    [0, 2],
    [2, 0],
    [2, 0],
    [2, 0],
    [2, 1],
    [2, 2],
    [2, 2],
    [2, 2],
    [2, 2],
    [2, 2],
    [2, 2]]
)
print('Raw data with true label and declared label')
print(data)

# For this problem, the confusion matrix can be computed directly: 
true_cm_count = np.zeros((4, 4), dtype=int)
for i in range(len(data)):
    true_cm_count[data[i, 0], data[i, 1]] += 1
print('True confusion matrix with counts')
print(true_cm_count)

# Normalization of the confusion matrix divides each row by the sum of that
# row; however, care must be taken for rows with no data.
row_sums = true_cm_count.sum(axis=1, keepdims=True)
row_divisor = np.maximum(row_sums, 1)
true_cm_norm = true_cm_count.astype(float) / row_divisor
print('Normalized confusion matrix by rows')
print(true_cm_norm)

print('')

target = torch.tensor(data[:, 0], device='cpu', dtype=torch.uint8)
pred = torch.tensor(data[:, 1], device='cpu', dtype=torch.uint8)

target_gpu = torch.tensor(data[:, 0], device='cuda', dtype=torch.uint8)
pred_gpu = torch.tensor(data[:, 1], device='cuda', dtype=torch.uint8)

print('Test 1: CPU tensors with no normalization')
metric_raw = pl.metrics.ConfusionMatrix(normalize=False)
cm_cpu_raw = metric_raw(pred, target)
print(cm_cpu_raw)
print('')

print('Test 2: CPU tensors with normalization')
metric_norm = pl.metrics.ConfusionMatrix(normalize=True)
cm_cpu_norm = metric_norm(pred, target)
print(cm_cpu_norm)
print('')

print('Test 3: GPU tensors with no normalization')
cm_gpu_raw = metric_raw(pred_gpu, target_gpu)
print(cm_gpu_raw)
print('')

print('Test 4: GPU tensors with normalization')
cm_gpu_norm = metric_norm(pred_gpu, target_gpu)
print(cm_gpu_norm)
print('')