import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class ChessPredictModelBaby(nn.Module):
	def __init__(self):
		super(ChessPredictModelBaby, self).__init__()
		self.model = nn.Sequential(
			nn.Conv2d(in_channels=8, out_channels=4, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(kernel_size=2, stride=2),

			nn.Conv2d(in_channels=4, out_channels=8, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.AdaptiveAvgPool2d(output_size=1),

			nn.Flatten(),
			nn.Linear(8, 16),
			nn.ReLU(),
			nn.Linear(16, 1),
			nn.Tanh()
		)

	def forward(self, x):
		x = self.model(x.permute(0, 3, 1, 2))
		return x

class ChessPredictModelS(nn.Module):
	def __init__(self):
		super(ChessPredictModelS, self).__init__()
		self.model = nn.Sequential(
			nn.Conv2d(in_channels=8, out_channels=32, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(kernel_size=2, stride=2),

			nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(kernel_size=2, stride=2),

			nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.AdaptiveAvgPool2d(output_size=1),

			nn.Flatten(),
			nn.Linear(128, 64),
			nn.ReLU(),
			nn.Linear(64, 1),
			nn.Tanh()
		)

	def forward(self, x):
		x = self.model(x.permute(0, 3, 1, 2))
		return x

class BasicBlock(nn.Module):
	def __init__(self, in_channels, out_channels, stride=1):
		super(BasicBlock, self).__init__()
		self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
		self.bn1 = nn.BatchNorm2d(out_channels)
		self.relu = nn.ReLU(inplace=True)
		self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
		self.bn2 = nn.BatchNorm2d(out_channels)

		self.downsample = None
		if stride != 1 or in_channels != out_channels:
			self.downsample = nn.Sequential(
				nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
				nn.BatchNorm2d(out_channels)
			)

	def forward(self, x):
		identity = x

		out = self.conv1(x)
		out = self.bn1(out)
		out = self.relu(out)

		out = self.conv2(out)
		out = self.bn2(out)

		if self.downsample is not None:
			identity = self.downsample(x)

		out += identity
		out = self.relu(out)

		return out

class ResNet(nn.Module):
	def __init__(self, block, layers, num_classes=1):
		super(ResNet, self).__init__()
		self.in_channels = 64

		self.model = nn.Sequential(
			nn.Conv2d(8, 64, kernel_size=3, stride=1, padding=1),
			nn.BatchNorm2d(64),
			nn.ReLU(inplace=True),
			self._make_layer(block, 64, layers[0], stride=1),
			self._make_layer(block, 128, layers[1], stride=2),
			self._make_layer(block, 256, layers[2], stride=2),
			self._make_layer(block, 512, layers[3], stride=2),
			nn.AdaptiveAvgPool2d((1, 1)),
			nn.Flatten(),
			nn.Linear(512, num_classes),
			nn.Tanh()
		)

	def _make_layer(self, block, out_channels, blocks, stride=1):
		layers = []
		layers.append(block(self.in_channels, out_channels, stride))
		self.in_channels = out_channels
		for _ in range(1, blocks):
			layers.append(block(out_channels, out_channels))
		return nn.Sequential(*layers)

	def forward(self, x):
		return self.model(x.permute(0, 3, 1, 2))

def resnet18():
	return ResNet(BasicBlock, [2, 2, 2, 2])

if __name__ == '__main__':
	model = ChessPredictModelS()
	summary(model, (8, 8, 8))