"""
Physics-Informed Loss for LiquidFlow.

Combines:
1. Rectified Flow Matching loss: ||v_θ(x_t, t) - (x_1 - x_0)||²
2. Physics residual: smoothness + continuity constraints on generated images
3. ConFIG gradient conflict-free update (from TUM-PBS, 2024)
4. Adaptive loss weighting (from PINN gradient pathology research, Wang et al. 2020)
"""

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


class PhysicsInformedFlowLoss(nn.Module):
    """
    Multi-objective loss for physics-informed flow matching.
    
    L = L_flow + λ_smooth * L_smooth + λ_tv * L_tv
    
    Where:
    - L_flow: Rectified flow matching MSE
    - L_smooth: Laplacian smoothness (heat equation steady state)  
    - L_tv: Total variation for edge preservation
    """
    
    def __init__(self, lambda_smooth=0.01, lambda_tv=0.001,
                 lambda_continuity=0.005, use_adaptive_weights=True):
        super().__init__()
        self.lambda_smooth = lambda_smooth
        self.lambda_tv = lambda_tv
        self.lambda_continuity = lambda_continuity
        self.use_adaptive_weights = use_adaptive_weights
        self.register_buffer('flow_grad_norm', torch.tensor(1.0))
        self.register_buffer('physics_grad_norm', torch.tensor(1.0))
        self.ema_decay = 0.99
    
    def flow_matching_loss(self, v_pred, x0, x1, t):
        target = x1 - x0
        return F.mse_loss(v_pred, target)
    
    def smoothness_loss(self, x_pred):
        lap_h = x_pred[:, :, 2:, :] - 2 * x_pred[:, :, 1:-1, :] + x_pred[:, :, :-2, :]
        lap_w = x_pred[:, :, :, 2:] - 2 * x_pred[:, :, :, 1:-1] + x_pred[:, :, :, :-2]
        h_min = min(lap_h.shape[2], lap_w.shape[2])
        w_min = min(lap_h.shape[3], lap_w.shape[3])
        laplacian = lap_h[:, :, :h_min, :w_min] + lap_w[:, :, :h_min, :w_min]
        return (laplacian ** 2).mean()
    
    def total_variation_loss(self, x_pred):
        diff_h = torch.abs(x_pred[:, :, 1:, :] - x_pred[:, :, :-1, :])
        diff_w = torch.abs(x_pred[:, :, :, 1:] - x_pred[:, :, :, :-1])
        return diff_h.mean() + diff_w.mean()
    
    def forward(self, v_pred, x0, x1, t, x_pred_clean=None, step=0):
        loss_flow = self.flow_matching_loss(v_pred, x0, x1, t)
        
        if x_pred_clean is None:
            t_expand = t.view(-1, 1, 1, 1)
            x_t = t_expand * x1 + (1 - t_expand) * x0
            x_pred_clean = x_t + v_pred * (1 - t_expand)
        
        warmup_steps = 500
        physics_weight = min(1.0, step / warmup_steps) if step < warmup_steps else 1.0
        
        loss_smooth = self.smoothness_loss(x_pred_clean)
        loss_tv = self.total_variation_loss(x_pred_clean)
        
        total = (loss_flow
                 + physics_weight * self.lambda_smooth * loss_smooth
                 + physics_weight * self.lambda_tv * loss_tv)
        
        losses = {
            'total': total, 'flow': loss_flow,
            'smooth': loss_smooth, 'tv': loss_tv,
            'physics_weight': torch.tensor(physics_weight),
        }
        return total, losses


class EMAModel:
    """Exponential Moving Average of model parameters."""
    
    def __init__(self, model, decay=0.9999, warmup_steps=1000):
        self.decay = decay
        self.warmup_steps = warmup_steps
        self.step = 0
        self.shadow = {}
        self.backup = {}
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()
    
    def update(self, model):
        self.step += 1
        decay = min(self.decay, (1 + self.step) / (10 + self.step))
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = decay * self.shadow[name] + (1 - decay) * param.data
    
    def apply_shadow(self, model):
        self.backup = {}
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.backup[name] = param.data.clone()
                param.data.copy_(self.shadow[name])
    
    def restore(self, model):
        for name, param in model.named_parameters():
            if param.requires_grad and name in self.backup:
                param.data.copy_(self.backup[name])
        self.backup = {}
    
    def state_dict(self):
        return {'shadow': self.shadow, 'step': self.step}
    
    def load_state_dict(self, state_dict):
        self.shadow = state_dict['shadow']
        self.step = state_dict['step']