# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03_loss.ipynb (unless otherwise specified).

__all__ = ['gram', 'content_loss', 'style_loss', 'tv_loss', 'pixel_loss', 'FastStyleLoss', 'SuperResLoss']

# Cell
from fastai.basics import *
from faststyle import *

# Cell
def gram(x):
    n,c,h,w = x.size()
    x = x.view(n, c, -1)
    return (x @ x.transpose(1,2))/(c*h*w)

# Cell
def content_loss(pred, targ, fts, layer_ws=None):
  layer_ws = layer_ws or [1.]*len(fts['pred']['cnt'])
  return sum([w*F.mse_loss(p,t) for w,p,t in zip_safe(layer_ws,fts['pred']['cnt'],fts['targ']['cnt'])])

# Cell
# TODO: Refactor
def style_loss(pred, targ, fts, layer_ws=None, stl_ws=None):
  bs = fts['pred']['stl'][0].shape[0]
  try:
    ws = fts['ws'].T[...,None,None]
  except KeyError:
    n = len(fts['source']['stl'][0])
    ws = torch.ones((n,bs,1,1)).to(default_device()) / n
  gs = []
  for fs in fts['source']['stl']:
    g = torch.stack([gram(f).expand(bs,-1,-1) for f in fs])
    g = g.mul_(ws).sum(dim=0)
    gs.append(g)
#   sims_gs = L(torch.stack([gram(f).expand(bs,-1,-1) for f in fs]) for fs in fts['source']['stl'])
#   sims_gs = L(torch.stack([gram(f) for f in fs]).wmean(stl_ws, dim=0) for fs in fts['source']['stl'])
  pred_gs = [gram(f) for f in fts['pred']['stl']]
  layer_ws = layer_ws or [1.]*len(pred_gs)
  # TODO: Use expand instead of repeat. Repeat copies the tensor data
  stl_losses = [w*F.mse_loss(g1,g2) for w,g1,g2 in zip_safe(layer_ws,gs,pred_gs)]
  return sum(stl_losses)

# Cell
def tv_loss(pred,*_):
  tv_h = ((pred[:,:,1:,:] - pred[:,:,:-1,:]).pow(2)).mean()
  tv_w = ((pred[:,:,:,1:] - pred[:,:,:,:-1]).pow(2)).mean()
  return (tv_h + tv_w)

# Cell
def pixel_loss(pred, targ, fts, base_loss=F.mse_loss, **kwargs):
  return base_loss(pred,targ,**kwargs)

# Cell
@funcs_kwargs
class FastStyleLoss:
  _methods=['stl_loss','cnt_loss','tv_loss']
  def __init__(self, stl_w=3e5, cnt_w=1, tv_w=20, **kwargs):
    store_attr('stl_w,cnt_w,tv_w')
    self.metrics = L(LossMetrics(['stl', 'cnt', 'tv']))

  def __call__(self, pred, targ, fts, **kwargs):
    self.stl = self.stl_w*self.stl_loss(pred, targ, fts)
    self.cnt = self.cnt_w*self.cnt_loss(pred, targ, fts)
    self.tv  = self.tv_w *self.tv_loss (pred, targ, fts)
    return self.stl+self.cnt+self.tv

  @delegates(style_loss)
  def stl_loss(self,pred,targ,fts,**kwargs): return style_loss(pred,targ,fts,**kwargs)
  @delegates(content_loss)
  def cnt_loss(self,pred,targ,fts,**kwargs): return content_loss(pred,targ,fts,**kwargs)
  @delegates(tv_loss)
  def tv_loss(self,pred,targ,fts,**kwargs): return tv_loss(pred,targ,fts,**kwargs)

# Cell
@funcs_kwargs
class SuperResLoss:
  _methods=['pixel_loss','cnt_loss']
  def __init__(self, **kwargs):
    self.metrics = L(LossMetrics(['pixel','cnt']))

  def __call__(self, pred, targ, fts, **kwargs):
    self.pixel = self.pixel_loss(pred,targ,fts,**kwargs)
    self.cnt   = self.cnt_loss  (pred,targ,fts,**kwargs)
    return self.pixel+self.cnt

  @delegates(pixel_loss)
  def pixel_loss(self,pred,targ,fts,**kwargs): return pixel_loss(pred,targ,fts,**kwargs)
  @delegates(content_loss)
  def cnt_loss(self,pred,targ,fts,**kwargs): return content_loss(pred,targ,fts,**kwargs)