import cv2
import torch
import torch.nn as nn

SIZE = 8


# Define a simple CNN using Sequential
class ChannelSum(nn.Module):
	def __init__(self):
		super(ChannelSum, self).__init__()
		self.model = nn.Sequential(
			nn.Conv2d(3, 3, kernel_size=SIZE),
			nn.Flatten(),
		)

	def forward(self, x):
		return self.model(x)


def main():
	model = ChannelSum()
	with torch.no_grad():
		for m in model.modules():
			if isinstance(m, nn.Conv2d):
				nn.init.constant_(m.bias, 0)
				m.bias[2] = -(SIZE * SIZE - 1)
				nn.init.constant_(m.weight, 0)
				# m.weight.shape = (out_channels, in_channels, kernel_size, kernel_size)
				# this essentially sums up the channel values
				for i in range(m.weight.shape[0]):
					m.weight[i, i, :, :] = 1

	print("running")
	model.eval()
	blue_only = torch.stack(
		[
			torch.zeros(1, 1, SIZE, SIZE),
			torch.zeros(1, 1, SIZE, SIZE),
			torch.ones(1, 1, SIZE, SIZE),
		]
	).reshape(1, 3, SIZE, SIZE)

	print("Expected results:", model(blue_only))
	print(
		"Channel order mismatch results:",
		model(blue_only.permute([0, 2, 3, 1]).reshape(1, 3, SIZE, SIZE)),
	)

	print("saving")
	torch.save(model.state_dict(), "weights.pth")
	cv2.imwrite(
		"testImage.png",
		blue_only.permute([0, 2, 3, 1]).squeeze().numpy()[:, :, ::-1] * 256,
	)


if __name__ == "__main__":
	main()
