import math
from types import SimpleNamespace
from typing import Tuple, Any

import torch
import torch.nn as nn
import torch.nn.functional as F

import yaml

# ----------------------------
# Misc helpers
# ----------------------------
def set_seed(seed: int = 42):
    import random
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


class AvgMeter:
    def __init__(self):
        self.reset()
    def reset(self):
        self.sum = 0.0
        self.cnt = 0
    def update(self, val, n=1):
        self.sum += float(val) * n
        self.cnt += n
    @property
    def avg(self):
        return self.sum / max(1, self.cnt)


# ----------------------------
# Image I/O helpers
# ----------------------------
def clamp_img(x: torch.Tensor) -> torch.Tensor:
    return torch.clamp(x, 0.0, 1.0)

def to_uint8(x: torch.Tensor) -> torch.Tensor:
    x = clamp_img(x)
    return (x * 255.0 + 0.5).to(torch.uint8)

def denorm(x: torch.Tensor, mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5)) -> torch.Tensor:
    """[-1,1] -> [0,1] if mean=std=0.5"""
    m = torch.tensor(mean, device=x.device).view(1, -1, 1, 1)
    s = torch.tensor(std, device=x.device).view(1, -1, 1, 1)
    return torch.clamp(x * s + m, 0.0, 1.0)

def norm(x: torch.Tensor, mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5)) -> torch.Tensor:
    m = torch.tensor(mean, device=x.device).view(1, -1, 1, 1)
    s = torch.tensor(std, device=x.device).view(1, -1, 1, 1)
    return (x - m) / s


# ----------------------------
# PSNR / SSIM (torch, no deps)
# ----------------------------
def psnr(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """pred/target in [0,1]. Returns scalar tensor (dB)."""
    mse = torch.mean((pred - target) ** 2)
    return 10.0 * torch.log10(1.0 / (mse + eps))

def _gaussian_window(window_size: int, sigma: float, device, dtype):
    coords = torch.arange(window_size, dtype=dtype, device=device) - (window_size - 1)/2
    g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
    g = g / g.sum()
    window = torch.outer(g, g)
    return window

def ssim(pred: torch.Tensor, target: torch.Tensor, window_size: int = 11, sigma: float = 1.5, eps: float = 1e-8) -> torch.Tensor:
    """
    pred/target: shape [B, C, H, W], values in [0,1].
    Returns scalar tensor averaged over batch.
    """
    assert pred.shape == target.shape
    B, C, H, W = pred.shape
    device, dtype = pred.device, pred.dtype

    window = _gaussian_window(window_size, sigma, device, dtype).view(1, 1, window_size, window_size)
    window = window.expand(C, 1, window_size, window_size)

    mu1 = F.conv2d(pred, window, groups=C, padding=window_size//2)
    mu2 = F.conv2d(target, window, groups=C, padding=window_size//2)

    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1 * mu2

    sigma1_sq = F.conv2d(pred * pred, window, groups=C, padding=window_size//2) - mu1_sq
    sigma2_sq = F.conv2d(target * target, window, groups=C, padding=window_size//2) - mu2_sq
    sigma12 = F.conv2d(pred * target, window, groups=C, padding=window_size//2) - mu1_mu2

    C1 = 0.01 ** 2
    C2 = 0.03 ** 2

    num = (2 * mu1_mu2 + C1) * (2 * sigma12 + C2)
    den = (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)

    ssim_map = num / (den + eps)
    return ssim_map.mean()


# ----------------------------
# EMA
# ----------------------------
class EMA:
    def __init__(self, model, decay=0.999):
        self.decay = decay
        self.shadow = {}
        for n, p in model.named_parameters():
            if p.requires_grad:
                self.shadow[n] = p.clone().detach()
    def update(self, model):
        with torch.no_grad():
            for n, p in model.named_parameters():
                if p.requires_grad:
                    self.shadow[n].mul_(self.decay).add_(p, alpha=1.0 - self.decay)
    def copy_to(self, model):
        with torch.no_grad():
            for n, p in model.named_parameters():
                if p.requires_grad:
                    p.copy_(self.shadow[n])


# ----------------------------
# Config helpers
# ----------------------------
def _to_namespace(value: Any) -> Any:
    if isinstance(value, dict):
        return SimpleNamespace(**{k: _to_namespace(v) for k, v in value.items()})
    if isinstance(value, list):
        return [_to_namespace(v) for v in value]
    return value

def load_config(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle) or {}
    return _to_namespace(data)


def resolve_device(device: str | None) -> str:
    if device:
        return device
    return "cuda" if torch.cuda.is_available() else "cpu"
