import torch
import torch.nn as nn

class ResBlock(nn.Module):
	"""Linear block with residuals"""
	def __init__(self, ch):
		super().__init__()
		self.join = nn.ReLU()
		self.long = nn.Sequential(
			nn.Linear(ch, ch),
			nn.LeakyReLU(0.1),
			nn.Linear(ch, ch),
			nn.LeakyReLU(0.1),
			nn.Linear(ch, ch),
		)
	def forward(self, x):
		return self.join(self.long(x) + x)

class AestheticPredictorModel(nn.Module):
	"""
	Main predictor class. Original:
	https://github.com/city96/CityClassifiers/blob/main/model.py
	"""
	def __init__(self, features=768, hidden=1024):
		super().__init__()
		self.features = features
		self.hidden = hidden
		self.up = nn.Sequential(
			nn.Linear(self.features, self.hidden),
			ResBlock(ch=self.hidden),
		)
		self.down = nn.Sequential(
			nn.Linear(self.hidden, 128),
			nn.Linear(128, 64),
			nn.Dropout(0.1),
			nn.LeakyReLU(),
			nn.Linear(64, 32),
			nn.Linear(32,  1),
			nn.Tanh(),
		)
	def forward(self, x):
		y = self.up(x)
		z = self.down(y)
		return (z+1.0)/2.0