"""
-------------------------------------------------------------------------------
Author: Dan64
Date: 2025-04-13
version:
LastEditors: Dan64
LastEditTime: 2025-04-13
-------------------------------------------------------------------------------
Description:
-------------------------------------------------------------------------------
Vapoursynth porting of some avisynth functions:

ColorYUV(autogain=True) -> AutoGain(clip)
ColorYUV(autowhite=True) -> AutoWhite(clip)
Autolevels() -> AutoLevels(clip, method, limit, gridsize, weight)

"""
import vapoursynth as vs
core = vs.core
import functools
import math
import numpy as np
import cv2
from PIL import Image
from functools import partial

"""
------------------------------------------------------------------------------- 
Author: Dan64
------------------------------------------------------------------------------- 
Description: Simple conversion of ColorYUV autogain function to python
URL: https://github.com/AviSynth/AviSynthPlus/blob/master/avs_core/filters/color.cpp 
------------------------------------------------------------------------------- 
"""

def AutoGain(clip: vs.VideoNode) -> vs.VideoNode:
   """ AutoGain filter using OpenCV

            :param clip:                Clip to process (support only RGB24).
   """
   if clip.format.id != vs.RGB24:
      # clip not in RGB24 format, it will be converted
      if clip.format.color_family == vs.ColorFamily.YUV:
         rgb_clip = clip.resize.Bicubic(format=vs.RGB24, matrix_in_s="709", range_s="full",
                                    dither_type="error_diffusion")
      else:
         rgb_clip = clip.resize.Bicubic(format=vs.RGB24, range_s="full")
   else:
      rgb_clip = clip

   def frame_autogain(n, f, bits_per_pixel: int = 8):

      img = _frame_to_image(f)
      img_np = np.asarray(img)

      yuv = cv2.cvtColor(img_np, cv2.COLOR_RGB2YUV)
      yuv_new = cv2.cvtColor(img_np, cv2.COLOR_RGB2YUV)

      dY = yuv[:, :, 0]

      if bits_per_pixel <= 16:
         bits_per_pixel_for_freq = bits_per_pixel
      else:
         bits_per_pixel_for_freq = 16

      loose_max_limit = (235 + 1) << (bits_per_pixel_for_freq - 8)
      loose_min_limit = 16 << (bits_per_pixel_for_freq - 8)
      loose_max = dY.max()
      loose_min = dY.min()
      maxY = min(loose_max, loose_max_limit)
      minY = max(loose_min, loose_min_limit)

      y_range = maxY - minY

      y_offset = 0
      y_gain = 0
      if y_range > 0:
         scale = (loose_max_limit - loose_min_limit) / y_range
         y_offset = (loose_min_limit - scale * minY) / (1 << (bits_per_pixel_for_freq - 8))
         y_gain = (256 * (scale - 1.0))

      dY = (dY + y_offset) * (y_gain/256 + 1)

      yuv_new[:, :, 0] = dY.clip(min=0, max=255)

      img_new = Image.fromarray(cv2.cvtColor(yuv_new, cv2.COLOR_YUV2RGB))

      return _image_to_frame(img_new, f.copy())

   clip_rgb = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(frame_autogain,
                                                                    bits_per_pixel = clip.format.bits_per_sample))

   if clip.format.id != vs.RGB24:
      # convert the format for tweak to YUV 8bits
      clip_new = clip_rgb.resize.Bicubic(format=vs.YUV420P8, matrix_s="709", range_s="limited")
   else:
      clip_new = clip_rgb

   return clip_new

"""
------------------------------------------------------------------------------- 
Author: Dan64
------------------------------------------------------------------------------- 
Description: Autolevels (i.e., histogram equalization) implementation in python
             using OpenCV
URL: https://pyimagesearch.com/2021/02/01/opencv-histogram-equalization-and-adaptive-histogram-equalization-clahe/
URL: https://docs.opencv.org/4.x/d5/daf/tutorial_py_histogram_equalization.html
------------------------------------------------------------------------------- 
"""

def AutoLevels(clip: vs.VideoNode, method: int = 3, clip_limit : float = 1.0, gridsize: int = 8,
               weight: float = 0) -> vs.VideoNode:
   """Histogram equalization implementation using OpenCV

      :param clip:                Clip to process (support only RGB24).
      :param method:              Method used to perform the histogram equalization.
                                  Allowed values are:
                                      0 : Apply Contrast Limited Adaptive Histogram Equalization on Luma (default)
                                      1 : Apply simple Histogram Equalization on all RGB channels
                                      2 : Apply CLAHE on all RGB channels
                                      3 : method=0 and method=1 are merged with assigned weight (default)
                                      4 : Automatic brightness and contrast optimization with ScaleAbs
                                      5 : Automatic brightness and contrast optimization with saturation arithmetic
      :param clip_limit:          Threshold for contrast limiting, used by method=(0,4) (default=1.1)
      :param gridsize:            Size of grid for histogram equalization, used by method=0.
                                  Input image will be divided into equally sized rectangular tiles.
                                  GridSize defines the number of tiles in row and column (default=8)
      :param weight:              Merges filter clip with original clip using the specified weight. A zero weight means
                                  that the filter clip is returned unchanged.
                                  (default=0)


   """
   if clip.format.id != vs.RGB24:
      # clip not in RGB24 format, it will be converted
      if clip.format.color_family == vs.ColorFamily.YUV:
         rgb_clip = clip.resize.Bicubic(format=vs.RGB24, matrix_in_s="709", range_s="full",
                                        dither_type="error_diffusion")
      else:
         rgb_clip = clip.resize.Bicubic(format=vs.RGB24, range_s="full")
   else:
      rgb_clip = clip

   #  Contrast Limited Adaptive Histogram Equalization
   def frame_autolevels_CLAHE(n, f, limit : float = 2.0, gridsize: int = 8):

      img = _frame_to_image(f)
      img_np = np.asarray(img)

      yuv = cv2.cvtColor(img_np, cv2.COLOR_RGB2YUV)

      y_image, u_image, v_image = cv2.split(yuv)

      clahe = cv2.createCLAHE(clipLimit=limit, tileGridSize=(gridsize, gridsize))

      y_image_eq = clahe.apply(y_image)

      yuv[:, :, 0] = y_image_eq.clip(min=0, max=255).astype(int)

      img_new = Image.fromarray(cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB))

      return _image_to_frame(img_new, f.copy())

   def frame_autolevels_CLAHE_rgb(n, f, limit : float = 2.0, gridsize: int = 8, algo = 0):

      img = _frame_to_image(f)
      img_np = np.asarray(img)

      r_image, g_image, b_image = cv2.split(img_np)

      if algo == 0:
         clahe = cv2.createCLAHE(clipLimit=limit, tileGridSize=(gridsize, gridsize))

         r_image_eq = clahe.apply(r_image)
         g_image_eq = clahe.apply(g_image)
         b_image_eq = clahe.apply(b_image)
      else:
         r_image_eq = cv2.equalizeHist(r_image)
         g_image_eq = cv2.equalizeHist(g_image)
         b_image_eq = cv2.equalizeHist(b_image)

      img_new = cv2.merge((r_image_eq, g_image_eq, b_image_eq))

      return _image_to_frame(img_new, f.copy())

   # Automatic brightness and contrast optimization with optional histogram clipping
   # https://stackoverflow.com/questions/56905592/automatic-contrast-and-brightness-adjustment-of-a-color-photo-of-a-sheet-of-pape
   def autolevels_with_Scale(n, f, clip_hist_percent: float = 1.0, algo: int = 0):

      img = _frame_to_image(f)
      img_np = np.asarray(img)

      gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)

      # Calculate grayscale histogram
      hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
      hist_size = len(hist)

      # Calculate cumulative distribution from the histogram
      accumulator = []
      accumulator.append(float(hist[0]))
      for index in range(1, hist_size):
         accumulator.append(accumulator[index - 1] + float(hist[index]))

      # Locate points to clip
      maximum = accumulator[-1]
      clip_hist_percent *= (maximum / 100.0)
      clip_hist_percent /= 2.0

      # Locate left cut
      minimum_gray = 0
      while accumulator[minimum_gray] < clip_hist_percent:
         minimum_gray += 1

      # Locate right cut
      maximum_gray = hist_size - 1
      while accumulator[maximum_gray] >= (maximum - clip_hist_percent):
         maximum_gray -= 1

      # Calculate alpha and beta values
      alpha = 255 / (maximum_gray - minimum_gray)
      beta = -minimum_gray * alpha

      if algo == 0:
         img_np_new = cv2.convertScaleAbs(img_np, alpha=alpha, beta=beta)
      else:
         # Add bias and gain to an image with saturation arithmetics. Unlike cv2.convertScaleAbs,
         # it does not take an absolute value, which would lead to nonsensical results
         # e.g., a pixel at 44 with alpha = 3 and beta = -210 becomes 78 with OpenCV, when in fact it should become 0.
         img_np_new = (img_np * alpha + beta).clip(min=0, max=255).astype(np.uint8)

      img_new = Image.fromarray(img_np_new)

      return _image_to_frame(img_new, f.copy())

   if method == 0:
      clip_a = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(frame_autolevels_CLAHE,
                                                                     limit=clip_limit, gridsize=gridsize))
      clip_b = rgb_clip
   elif method == 1:
      clip_a = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(frame_autolevels_CLAHE_rgb, algo=1,
                                                                     limit=clip_limit, gridsize=gridsize))
      clip_b = rgb_clip
   elif method == 2:
      clip_a = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(frame_autolevels_CLAHE_rgb, algo=0,
                                                                     limit=clip_limit, gridsize=gridsize))
      clip_b = rgb_clip
   elif method == 3:
      clip_a = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(frame_autolevels_CLAHE,
                                                                       limit=clip_limit, gridsize=gridsize))
      clip_b = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(frame_autolevels_CLAHE_rgb, algo=1,
                                                                     limit=clip_limit, gridsize=gridsize))
   elif method ==4:
      clip_a = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(autolevels_with_Scale, algo=0,
                                                                     clip_hist_percent=clip_limit))
      clip_b = rgb_clip
   else:
      clip_a = clip.std.ModifyFrame(clips=rgb_clip, selector=partial(autolevels_with_Scale, algo=1,
                                                                     clip_hist_percent=clip_limit))
      clip_b = rgb_clip

   if weight > 0:
      clip_rgb = core.std.Merge(clip_a, clip_b, weight)
   else:
      clip_rgb = clip_a

   if clip.format.id != vs.RGB24:
      # convert the format for tweak to YUV 8bits
      clip_new = clip_rgb.resize.Bicubic(format=vs.YUV420P8, matrix_s="709", range_s="limited")
   else:
      clip_new = clip_rgb

   return clip_new

"""
------------------------------------------------------------------------------- 
Author: Dan64
------------------------------------------------------------------------------- 
Description: function that takes a video clip as an input and calculates the average color values 
for each of the three color planes (red, green, blue). The white balance of the input clip is adjusted based on the 
color balance of the individual frames.
URL: http://www.vapoursynth.com/doc/functions/frameeval.html
------------------------------------------------------------------------------- 
"""
def AutoWhite(clip: vs.VideoNode) -> vs.VideoNode:
   """ Autowhite filter using PlaneStats()

         :param clip:                Clip to process (support only RGB24).
   """
   if clip.format.id != vs.RGB24:
      # clip not in RGB24 format, it will be converted
      if clip.format.color_family == vs.ColorFamily.YUV:
         rgb_clip = clip.resize.Bicubic(format=vs.RGB24, matrix_in_s="709", range_s="full",
                                        dither_type="error_diffusion")
      else:
         rgb_clip = clip.resize.Bicubic(format=vs.RGB24, range_s="full")
   else:
      rgb_clip = clip

   # auto white from http://www.vapoursynth.com/doc/functions/frameeval.html
   def frame_autowhite(n, f, clip, core):
      small_number = 0.000000001
      red = f[0].props['PlaneStatsAverage']
      green = f[1].props['PlaneStatsAverage']
      blue = f[2].props['PlaneStatsAverage']
      max_rgb = max(red, green, blue)
      red_corr = max_rgb / max(red, small_number)
      green_corr = max_rgb / max(green, small_number)
      blue_corr = max_rgb / max(blue, small_number)
      norm = max(blue, math.sqrt(red_corr * red_corr + green_corr * green_corr + blue_corr * blue_corr) / math.sqrt(3),
                 small_number)
      r_gain = red_corr / norm
      g_gain = green_corr / norm
      b_gain = blue_corr / norm
      return core.std.Expr(clip,
                           expr=['x ' + repr(r_gain) + ' *', 'x ' + repr(g_gain) + ' *', 'x ' + repr(b_gain) + ' *'])

   r_avg = core.std.PlaneStats(rgb_clip, plane=0)
   g_avg = core.std.PlaneStats(rgb_clip, plane=1)
   b_avg = core.std.PlaneStats(rgb_clip, plane=2)

   clip_rgb = core.std.FrameEval(rgb_clip, functools.partial(frame_autowhite, clip=rgb_clip, core=vs.core),
                             prop_src=[r_avg, g_avg, b_avg])

   if clip.format.id != vs.RGB24:
      # convert the format for tweak to YUV 8bits
      clip_new = clip_rgb.resize.Bicubic(format=vs.YUV420P8, matrix_s="709", range_s="limited")
   else:
      clip_new = clip_rgb

   return clip_new

def _frame_to_image(frame: vs.VideoFrame) -> Image:
   npArray = np.dstack([np.asarray(frame[plane]) for plane in range(frame.format.num_planes)])
   return Image.fromarray(npArray, 'RGB')


def _image_to_frame(img: Image, frame: vs.VideoFrame) -> vs.VideoFrame:
   npArray = np.array(img)
   [np.copyto(np.asarray(frame[plane]), npArray[:, :, plane]) for plane in range(frame.format.num_planes)]
   return frame
