##  Copyright (c) 2020 The WebM project authors. All Rights Reserved.
##
##  Use of this source code is governed by a BSD-style license
##  that can be found in the LICENSE file in the root of the source
##  tree. An additional intellectual property rights grant can be found
##  in the file PATENTS.  All contributing project authors may
##  be found in the AUTHORS file in the root of the source tree.
##

# coding: utf-8
import numpy as np
import numpy.linalg as LA
from Util import MSE
from MotionEST import MotionEST
"""Search & Smooth Model with Adapt Weights"""


class SearchSmoothAdapt(MotionEST):
  """
    Constructor:
        cur_f: current frame
        ref_f: reference frame
        blk_sz: block size
        wnd_size: search window size
        beta: neigbor loss weight
        max_iter: maximum number of iterations
        metric: metric to compare the blocks distrotion
    """

  def __init__(self, cur_f, ref_f, blk_size, search, max_iter=100):
    self.search = search
    self.max_iter = max_iter
    super(SearchSmoothAdapt, self).__init__(cur_f, ref_f, blk_size)

  """
    get local diffiencial of refernce
    """

  def getRefLocalDiff(self, mvs):
    m, n = self.num_row, self.num_col
    localDiff = [[] for _ in xrange(m)]
    blk_sz = self.blk_sz
    for r in xrange(m):
      for c in xrange(n):
        I_row = 0
        I_col = 0
        #get ssd surface
        count = 0
        center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz,
                              c * blk_sz:(c + 1) * blk_sz, 0]
        ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz)
        tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz)
        target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0]
        for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}:
          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
            I_row += np.sum(np.abs(nb - center)) - np.sum(
                np.abs(target - center))
            count += 1
        I_row //= (count * blk_sz * blk_sz)
        count = 0
        for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}:
          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
            I_col += np.sum(np.abs(nb - center)) - np.sum(
                np.abs(target - center))
            count += 1
        I_col //= (count * blk_sz * blk_sz)
        localDiff[r].append(
            np.array([[I_row * I_row, I_row * I_col],
                      [I_col * I_row, I_col * I_col]]))
    return localDiff

  """
    add smooth constraint
    """

  def smooth(self, uvs, mvs):
    sm_uvs = np.zeros(uvs.shape)
    blk_sz = self.blk_sz
    for r in xrange(self.num_row):
      for c in xrange(self.num_col):
        nb_uv = np.array([0.0, 0.0])
        for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}:
          if 0 <= i < self.num_row and 0 <= j < self.num_col:
            nb_uv += uvs[i, j] / 6.0
          else:
            nb_uv += uvs[r, c] / 6.0
        for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1),
                     (r + 1, c + 1)}:
          if 0 <= i < self.num_row and 0 <= j < self.num_col:
            nb_uv += uvs[i, j] / 12.0
          else:
            nb_uv += uvs[r, c] / 12.0
        ssd_nb = self.block_dist(r, c, self.blk_sz * nb_uv)
        mv = mvs[r, c]
        ssd_mv = self.block_dist(r, c, mv)
        alpha = (ssd_nb - ssd_mv) / (ssd_mv + 1e-6)
        M = alpha * self.localDiff[r][c]
        P = M + np.identity(2)
        inv_P = LA.inv(P)
        sm_uvs[r, c] = np.dot(inv_P, nb_uv) + np.dot(
            np.matmul(inv_P, M), mv / blk_sz)
    return sm_uvs

  def block_matching(self):
    self.search.motion_field_estimation()

  def motion_field_estimation(self):
    self.localDiff = self.getRefLocalDiff(self.search.mf)
    #get matching results
    mvs = self.search.mf
    #add smoothness constraint
    uvs = mvs / self.blk_sz
    for _ in xrange(self.max_iter):
      uvs = self.smooth(uvs, mvs)
    self.mf = uvs * self.blk_sz


"""Search & Smooth Model with Fixed Weights"""


class SearchSmoothFix(MotionEST):
  """
    Constructor:
        cur_f: current frame
        ref_f: reference frame
        blk_sz: block size
        wnd_size: search window size
        beta: neigbor loss weight
        max_iter: maximum number of iterations
        metric: metric to compare the blocks distrotion
    """

  def __init__(self, cur_f, ref_f, blk_size, search, beta, max_iter=100):
    self.search = search
    self.max_iter = max_iter
    self.beta = beta
    super(SearchSmoothFix, self).__init__(cur_f, ref_f, blk_size)

  """
    get local diffiencial of refernce
    """

  def getRefLocalDiff(self, mvs):
    m, n = self.num_row, self.num_col
    localDiff = [[] for _ in xrange(m)]
    blk_sz = self.blk_sz
    for r in xrange(m):
      for c in xrange(n):
        I_row = 0
        I_col = 0
        #get ssd surface
        count = 0
        center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz,
                              c * blk_sz:(c + 1) * blk_sz, 0]
        ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz)
        tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz)
        target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0]
        for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}:
          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
            I_row += np.sum(np.abs(nb - center)) - np.sum(
                np.abs(target - center))
            count += 1
        I_row //= (count * blk_sz * blk_sz)
        count = 0
        for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}:
          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
            I_col += np.sum(np.abs(nb - center)) - np.sum(
                np.abs(target - center))
            count += 1
        I_col //= (count * blk_sz * blk_sz)
        localDiff[r].append(
            np.array([[I_row * I_row, I_row * I_col],
                      [I_col * I_row, I_col * I_col]]))
    return localDiff

  """
    add smooth constraint
    """

  def smooth(self, uvs, mvs):
    sm_uvs = np.zeros(uvs.shape)
    blk_sz = self.blk_sz
    for r in xrange(self.num_row):
      for c in xrange(self.num_col):
        nb_uv = np.array([0.0, 0.0])
        for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}:
          if 0 <= i < self.num_row and 0 <= j < self.num_col:
            nb_uv += uvs[i, j] / 6.0
          else:
            nb_uv += uvs[r, c] / 6.0
        for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1),
                     (r + 1, c + 1)}:
          if 0 <= i < self.num_row and 0 <= j < self.num_col:
            nb_uv += uvs[i, j] / 12.0
          else:
            nb_uv += uvs[r, c] / 12.0
        mv = mvs[r, c] / blk_sz
        M = self.localDiff[r][c]
        P = M + self.beta * np.identity(2)
        inv_P = LA.inv(P)
        sm_uvs[r, c] = np.dot(inv_P, self.beta * nb_uv) + np.dot(
            np.matmul(inv_P, M), mv)
    return sm_uvs

  def block_matching(self):
    self.search.motion_field_estimation()

  def motion_field_estimation(self):
    #get local structure
    self.localDiff = self.getRefLocalDiff(self.search.mf)
    #get matching results
    mvs = self.search.mf
    #add smoothness constraint
    uvs = mvs / self.blk_sz
    for _ in xrange(self.max_iter):
      uvs = self.smooth(uvs, mvs)
    self.mf = uvs * self.blk_sz