from __future__ import annotations

import torch
import torch.nn as nn

def timestep_embedding(timesteps: torch.Tensor, dim: int) -> torch.Tensor:
	if timesteps.ndim == 0:
		timesteps = timesteps.unsqueeze(0)
	half = dim // 2
	freqs = torch.exp(
		-torch.log(torch.tensor(10000.0, device=timesteps.device))
		* torch.arange(0, half, device=timesteps.device)
		/ max(1, half)
	)
	args = timesteps.float().unsqueeze(1) * freqs.unsqueeze(0)
	emb = torch.cat([torch.cos(args), torch.sin(args)], dim=1)
	if dim % 2 == 1:
		emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=1)
	return emb


class ResBlock(nn.Module):
	def __init__(self, channels: int, emb_dim: int):
		super().__init__()
		self.conv1 = nn.Conv2d(channels, channels, 3, 1, 1)
		self.conv2 = nn.Conv2d(channels, channels, 3, 1, 1)
		self.norm1 = nn.GroupNorm(8, channels)
		self.norm2 = nn.GroupNorm(8, channels)
		self.emb = nn.Linear(emb_dim, channels)
		self.act = nn.SiLU()

	def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
		h = self.norm1(x)
		h = h + self.emb(t_emb).unsqueeze(-1).unsqueeze(-1)
		h = self.act(h)
		h = self.conv1(h)
		h = self.norm2(h)
		h = self.act(h)
		h = self.conv2(h)
		return x + h


class SimpleUNet(nn.Module):
	def __init__(
		self,
		in_ch: int,
		base: int,
		emb_dim: int,
		num_variants: int,
		x_ch: int = 3,
		cond_ch: int | None = None,
	):
		super().__init__()
		self.x_ch = x_ch
		if cond_ch is None:
			cond_ch = in_ch - x_ch
		self.cond_ch = cond_ch
		if self.cond_ch <= 0:
			raise ValueError("cond_ch must be positive")
		total_in = x_ch + self.cond_ch
		if in_ch != total_in:
			raise ValueError(
				f"SimpleUNet configured with in_ch={in_ch}, but got x_ch={x_ch} and cond_ch={self.cond_ch}."
			)
		self.in_conv = nn.Conv2d(total_in, base, 3, 1, 1)
		self.down1 = nn.Conv2d(base, base * 2, 4, 2, 1)
		self.down2 = nn.Conv2d(base * 2, base * 4, 4, 2, 1)

		self.rb1 = ResBlock(base, emb_dim)
		self.rb2 = ResBlock(base * 2, emb_dim)
		self.rb3 = ResBlock(base * 4, emb_dim)

		self.up1 = nn.ConvTranspose2d(base * 4, base * 2, 4, 2, 1)
		self.up2 = nn.ConvTranspose2d(base * 2, base, 4, 2, 1)
		self.out = nn.Conv2d(base, 3, 3, 1, 1)

		self.variant_embed = nn.Embedding(num_variants, emb_dim)
		self.emb_mlp = nn.Sequential(
			nn.Linear(emb_dim, emb_dim),
			nn.SiLU(),
			nn.Linear(emb_dim, emb_dim),
		)

	def forward(
		self,
		x_t: torch.Tensor,
		t: torch.Tensor,
		y: torch.Tensor,
		variant_idx: torch.Tensor,
	) -> torch.Tensor:
		if x_t.size(1) != self.x_ch:
			raise ValueError(f"Expected {self.x_ch} diffusion channels, got {x_t.size(1)}")
		if y.size(1) != self.cond_ch:
			raise ValueError(f"Expected {self.cond_ch} conditioning channels, got {y.size(1)}")
		emb_dim = self.emb_mlp[0].in_features
		t_emb = timestep_embedding(t, dim=emb_dim).to(x_t.device)
		v_emb = self.variant_embed(variant_idx.to(x_t.device))
		cond_emb = self.emb_mlp(t_emb + v_emb)

		h0 = torch.cat([x_t, y], dim=1)
		h1 = self.in_conv(h0)
		h1 = self.rb1(h1, cond_emb)

		h2 = self.down1(h1)
		h2 = self.rb2(h2, cond_emb)

		h3 = self.down2(h2)
		h3 = self.rb3(h3, cond_emb)

		u2 = self.up1(h3)
		u1 = self.up2(u2)
		return self.out(u1)