1 files changed, 64 insertions, 34 deletions

diff --git a/lvc/src/encode.rs b/lvc/src/encode.rs
index 10f8c50..fde4f94 100644
--- a/lvc/src/encode.rs
+++ b/lvc/src/encode.rs
@@ -1,6 +1,6 @@
 use crate::diff::{diff, pixel_diff};
 use crate::split::split;
-use crate::{Block, Frame, Ref, View, P2};
+use crate::{Block, Frame, Pixel, Ref, View, P2};
 
 pub struct EncodeConfig {
     pub threshold: u32,
@@ -25,15 +25,15 @@ pub fn encode(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConf
     let att = 1. - attention(frame, view) as f32 * 0.000001;
     let thres = (config.threshold as f32 * att.clamp(0.2, 1.0)) as u32;
 
-    for granularity in [4, 2, 1] {
-        for _ in 0..config.iters {
-            let (nd, nrp) = optimize_ref(last_frame, frame, view, d, r, granularity);
-            if nd < d {
-                r = nrp;
-                d = nd;
-            } else {
-                break;
-            }
+    let target_average = average_color(frame, view);
+
+    for granularity in [8, 8, 4, 2, 1, 1, 1] {
+        let (nd, nrp) = optimize_ref(last_frame, frame, view, r, granularity, target_average);
+        if nd < d {
+            r = nrp;
+            d = nd;
+        } else {
+            break;
         }
     }
 
@@ -48,33 +48,42 @@ pub fn optimize_ref(
     last_frame: &Frame,
     frame: &Frame,
     view: View,
-    mut d: u32,
-    mut r: Ref,
+    r: Ref,
     granularity: i32,
+    target_average: Pixel,
 ) -> (u32, Ref) {
-    let mut n = |f: fn(&mut Ref, i32)| {
-        let mut r2 = r;
-        f(&mut r2, granularity);
-
-        let d2 = diff([last_frame, frame], view, r2);
-        if d2 < d {
-            d = d2;
-            r = r2;
-        }
-    };
-
-    n(|r, g| r.pos_off.x += g);
-    n(|r, g| r.pos_off.x -= g);
-    n(|r, g| r.pos_off.y += g);
-    n(|r, g| r.pos_off.y -= g);
-    n(|r, g| r.color_off.r += (g as i16) << 2);
-    n(|r, g| r.color_off.r -= (g as i16) << 2);
-    n(|r, g| r.color_off.g += (g as i16) << 2);
-    n(|r, g| r.color_off.g -= (g as i16) << 2);
-    n(|r, g| r.color_off.b += (g as i16) << 2);
-    n(|r, g| r.color_off.b -= (g as i16) << 2);
+    [
+        Some(r.apply(|r| r.pos_off.x += granularity)),
+        Some(r.apply(|r| r.pos_off.x -= granularity)),
+        Some(r.apply(|r| r.pos_off.y += granularity)),
+        Some(r.apply(|r| r.pos_off.y -= granularity)),
+        {
+            let mut r = r;
+            let last_avr = average_color(last_frame, view);
+            let diff = target_average - last_avr;
+            r.color_off = diff;
+            if diff != Pixel::BLACK {
+                Some(r)
+            } else {
+                None
+            }
+        },
+        // n(|r, g| r.color_off.r += (g as i16) << 2);
+        // n(|r, g| r.color_off.r -= (g as i16) << 2);
+        // n(|r, g| r.color_off.g += (g as i16) << 2);
+        // n(|r, g| r.color_off.g -= (g as i16) << 2);
+        // n(|r, g| r.color_off.b += (g as i16) << 2);
+        // n(|r, g| r.color_off.b -= (g as i16) << 2);
+    ]
+    .into_iter()
+    .flatten()
+    .map(|r| {
+        let d = diff([last_frame, frame], view, r);
 
-    (d, r)
+        (d, r)
+    })
+    .min_by_key(|e| e.0)
+    .unwrap()
 }
 
 pub fn attention(frame: &Frame, view: View) -> u32 {
@@ -88,3 +97,24 @@ pub fn attention(frame: &Frame, view: View) -> u32 {
     }
     k
 }
+
+pub fn average_color(frame: &Frame, view: View) -> Pixel {
+    let mut r = 0u32;
+    let mut g = 0u32;
+    let mut b = 0u32;
+
+    for y in view.a.y..view.b.y {
+        for x in view.a.x..view.b.x {
+            let p = frame[P2 { x, y }];
+            r += p.r as u32;
+            g += p.g as u32;
+            b += p.b as u32;
+        }
+    }
+    let area = view.size().area() as u32;
+    Pixel {
+        r: (r / area) as i16,
+        g: (g / area) as i16,
+        b: (b / area) as i16,
+    }
+}