"""Metrics are fidelity/artifact proxies; none certifies plausible color by itself."""
import numpy as np
import torch

def soft_ce(logits, idx, weights, class_weights=None):
    # FP32 reduction under autocast; no full target distribution materialized.
    logits=logits.float()
    idx=idx.permute(0,3,1,2)
    weights=weights.permute(0,3,1,2)
    nll=-((logits.gather(1,idx)-logits.logsumexp(1,keepdim=True))*weights).sum(1)
    if class_weights is not None:
        nll=nll*class_weights[idx[:,0]]
    return nll.mean()

def per_image_metrics(L, pred, target):
    error=(pred-target).square().sum(1).sqrt().mean((1,2))
    chroma=pred.square().sum(1).sqrt().mean((1,2))
    true_chroma=target.square().sum(1).sqrt().mean((1,2))
    seams=[]
    for dim in [-1,-2]:
        dp=pred.diff(dim=dim).square().sum(1).sqrt()
        dt=target.diff(dim=dim).square().sum(1).sqrt()
        dl=L.diff(dim=dim).abs()[:,0]*50
        # Penalize excess chroma discontinuity where luminance is flat.
        mask=(dl<2).float()
        seams.append(((dp-dt).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
    return {'ab_error':error,'chroma':chroma,'target_chroma':true_chroma,
            'excess_chroma_edge':sum(seams)/2}

def summarize(rows):
    keys=['ce','ab_error','chroma','target_chroma','excess_chroma_edge']
    result={k:float(np.mean([r[k] for r in rows])) for k in keys}
    result['n']=len(rows)
    result['chroma_ratio']=result['chroma']/max(result['target_chroma'],1e-9)
    return result