from __future__ import annotations

from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, Dict

import torch
import torch.nn as nn

from attention_net import AttentionNet
from ddpm_unet import DiffusionUNet
from simple_unet import SimpleUNet

@dataclass
class DiffusionConfig:
	T: int = 1000
	beta_start: float = 1e-4
	beta_end: float = 0.02


class RainDiffusionModel(nn.Module):
	def __init__(
		self,
		unet: nn.Module,
		cfg: DiffusionConfig,
		attention_net: nn.Module | None = None,
	):
		super().__init__()
		self.unet = unet
		self.cfg = cfg
		self.attention_net = attention_net
		self.use_attention = attention_net is not None

		betas = torch.linspace(cfg.beta_start, cfg.beta_end, cfg.T)
		alphas = 1.0 - betas
		self.register_buffer("betas", betas)
		self.register_buffer("alphas", alphas)
		self.register_buffer("alphas_cumprod", torch.cumprod(alphas, dim=0))
		self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(self.alphas_cumprod))
		self.register_buffer(
			"sqrt_one_minus_alphas_cumprod",
			torch.sqrt(1.0 - self.alphas_cumprod),
		)

	def q_sample(
		self,
		x0: torch.Tensor,
		t: torch.Tensor,
		noise: torch.Tensor | None = None,
	) -> tuple[torch.Tensor, torch.Tensor]:
		if noise is None:
			noise = torch.randn_like(x0)
		a = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
		b = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
		return a * x0 + b * noise, noise

	def predict_eps(
		self,
		x_t: torch.Tensor,
		t: torch.Tensor,
		y: torch.Tensor,
		variant_idx: torch.Tensor | None,
	) -> torch.Tensor:
		if not self.use_attention:
			return self.unet(x_t, t, y, variant_idx)

		m = self.attention_net(y)
		cond_y = torch.cat([y, m], dim=1)
		return self.unet(x_t, t, cond_y, variant_idx)

	@torch.no_grad()
	def p_sample(
		self,
		x: torch.Tensor,
		t: torch.Tensor,
		y: torch.Tensor,
		variant_idx: torch.Tensor | None,
	) -> torch.Tensor:
		eps = self.predict_eps(x, t, y, variant_idx)
		alpha_bar_t = self.alphas_cumprod[t].view(-1, 1, 1, 1)
		x0_hat = (x - torch.sqrt(1.0 - alpha_bar_t) * eps) / torch.sqrt(alpha_bar_t)
		x0_hat = torch.clamp(x0_hat, 0.0, 1.0)

		if t[0].item() == 0:
			return x0_hat

		t_prev = torch.clamp(t - 1, min=0)
		alpha_bar_prev = self.alphas_cumprod[t_prev].view(-1, 1, 1, 1)
		noise = torch.randn_like(x)
		x_prev = torch.sqrt(alpha_bar_prev) * x0_hat + torch.sqrt(1.0 - alpha_bar_prev) * noise
		return x_prev

	@torch.no_grad()
	def sample(self, y: torch.Tensor, variant_idx: torch.Tensor | None) -> torch.Tensor:
		device = y.device
		x = torch.randn_like(y)
		for t_int in reversed(range(self.cfg.T)):
			t = torch.full((y.size(0),), t_int, device=device, dtype=torch.long)
			x = self.p_sample(x, t, y, variant_idx)
		return torch.clamp(x, 0.0, 1.0)


def _get_config_dict(cfg_obj: Any, defaults: Dict = None) -> Dict:
    defaults = defaults or {}
    return {key: getattr(cfg_obj, key, default) for key, default in defaults.items()}


def build_model(cfg: Any, device: str = "cuda") -> Any:
	"""
	Build the rain removal diffusion model with specified backbone.
	
	Args:
		cfg: Configuration object with model and data settings
		device: Device to load model on
	
	Returns:
		RainDiffusionModel instance
	"""
	model_cfg = cfg.model
	backbone = getattr(model_cfg, "backbone", "simple_unet").lower()

	attention_cfg = getattr(model_cfg, "attention", SimpleNamespace())
	attention_enabled = getattr(attention_cfg, "enabled", False)
	extra_cond_ch = 1 if attention_enabled else 0

	attention_net: AttentionNet | None = None
	if attention_enabled:
		att_in_ch = getattr(attention_cfg, "in_channels", 3)
		att_base = getattr(attention_cfg, "base_channels", 32)
		att_max = getattr(attention_cfg, "max_channels", 128)
		attention_net = AttentionNet(
			in_channels=att_in_ch,
			base_channels=att_base,
			max_channels=att_max,
		)
		pretrained_path = getattr(attention_cfg, "pretrained_path", None)
		if pretrained_path:
			state = torch.load(pretrained_path, map_location=device)
			if isinstance(state, dict) and "state_dict" in state:
				state = state["state_dict"]
			attention_net.load_state_dict(state, strict=False)
		attention_net = attention_net.to(device)

	if backbone == "simple_unet":
		simple_cfg = getattr(model_cfg, "simple_unet", SimpleNamespace())
		num_variants = getattr(simple_cfg, "num_variants", getattr(cfg.data, "num_variants", 3))
		x_ch = getattr(simple_cfg, "x_ch", 3)
		base_cond = getattr(simple_cfg, "cond_ch", max(0, getattr(simple_cfg, "in_ch", 6) - x_ch))
		cond_ch = base_cond + extra_cond_ch
		total_in = x_ch + cond_ch

		net = SimpleUNet(
			in_ch=total_in,
			base=getattr(simple_cfg, "base", 64),
			emb_dim=getattr(simple_cfg, "emb_dim", 256),
			num_variants=num_variants,
			x_ch=x_ch,
			cond_ch=cond_ch,
		)

	elif backbone == "ddpm_unet":
		ddpm_cfg = getattr(model_cfg, "ddpm_unet", SimpleNamespace())
		defaults = {
			"in_ch": 3,
			"cond_ch": 3,
			"out_ch": 3,
			"base_ch": 128,
			"channel_mults": (1, 1, 2, 2, 4, 4),
			"num_res_blocks": 2,
			"attn_resolutions": (16,),
			"dropout": 0.0,
			"resamp_with_conv": True,
			"num_variants": getattr(cfg.data, "num_variants", 3),
			"image_size": getattr(cfg.data, "image_size", 256),
		}

		unet_config = _get_config_dict(ddpm_cfg, defaults)
		if isinstance(unet_config.get("channel_mults"), list):
			unet_config["channel_mults"] = tuple(unet_config["channel_mults"])
		if isinstance(unet_config.get("attn_resolutions"), list):
			unet_config["attn_resolutions"] = tuple(unet_config["attn_resolutions"])
		base_cond = unet_config.get("cond_ch", defaults["cond_ch"])
		unet_config["cond_ch"] = base_cond + extra_cond_ch
		net = DiffusionUNet(**unet_config)
	else:
		raise ValueError(f"Unknown backbone: {backbone}. Choose 'simple_unet' or 'ddpm_unet'")

	# Setup diffusion configuration
	diffusion_cfg = getattr(model_cfg, "diffusion", None)
	if diffusion_cfg is None:
		diff_cfg = DiffusionConfig()
	else:
		diff_defaults = {'T': 1000, 'beta_start': 1e-4, 'beta_end': 0.02}
		diff_params = _get_config_dict(diffusion_cfg, diff_defaults)
		diff_cfg = DiffusionConfig(**diff_params)

	return RainDiffusionModel(net.to(device), diff_cfg, attention_net=attention_net).to(device)
