import numpy as np
import torch
import torch.nn as nn
import math
from torchinfo import summary

# https://arxiv.org/abs/1502.03167
# https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html

# 1. Normalizace celého datasetu (před vstupem do sítě):
# Např. převod [0, 255] na [0, 1] nebo standardizace přes transforms.Normalize(mean, std):
# obraz = (obraz - průměr) / směrodatná_odchylka

# 2. Dávková normalizace (Batch Normalization) uvnitř sítě:
# Průměr a rozptyl se počítají z aktuální minidávky.
# Hodnoty se normalizují a následně upraví trénovatelnými parametry gamma a beta (měřítko a posun):
# y = gamma * ((x - průměr) / sqrt(rozptyl + eps)) + beta

# --- Ruční výpočet na minidávce 3 čísel ---
data = [3, 6, 18]
print("Minidávka 3 hodnot:", data)

mini_batch_mean = sum(data) / len(data)
print("Průměr dávky:", mini_batch_mean)

mini_batch_variance = sum((xi - mini_batch_mean) ** 2 for xi in data) / len(data)
print("Rozptyl dávky:", mini_batch_variance)

normalize = []
for xi in data:
    num = xi - mini_batch_mean
    den = math.sqrt(mini_batch_variance + 0.0001)
    normalize.append(num / den)
print("Ručně normalizované hodnoty:", normalize)


# --- Stejný výpočet pomocí PyTorch BatchNorm1d ---
# Tenzor o rozměru [batch_size=1, channels=1, délka=3]
data_tensor = torch.tensor([3.0, 6.0, 18.0]).unsqueeze(0).unsqueeze(0)
bn = nn.BatchNorm1d(1, eps=0.0001)
output = bn(data_tensor)
print("Výstup z PyTorch BatchNorm1d:", output)


# --- Příklad 2D dávkové normalizace pro obraz: BatchNorm2d ---
# Typické umístění: mezi konvoluci a aktivační funkci (ReLU).
model = torch.nn.Sequential(
    torch.nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1),
    torch.nn.BatchNorm2d(6),  # počet kanálů musí odpovídat out_channels z Conv2d
    torch.nn.ReLU(),
    torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
)

test_input = torch.randn(20, 1, 32, 32)
output = model(test_input)

# BatchNorm2d(6) má 12 parametrů:
# Pro každý ze 6 kanálů má 2 trénovatelné parametry: gamma (měřítko) a beta (posun) -> 6 * 2 = 12
summary(model, input_size=(20, 1, 32, 32), col_names=("input_size", "output_size", "num_params"), device="cpu")

# model.eval() zajistí, že se průměr a rozptyl nepočítají z testovacích dat,
# ale použijí se hodnoty nasbírané při trénování.
model.eval()