]> git.sesse.net Git - narabu/blobdiff - qdc.cpp
Silence some Mesa warnings.
[narabu] / qdc.cpp
diff --git a/qdc.cpp b/qdc.cpp
index 43fcba0de2dc343fbdee87d762992fa00c9ff95b..0339e12af0ad07ba2616530b30bd1b2012258745 100644 (file)
--- a/qdc.cpp
+++ b/qdc.cpp
@@ -1,6 +1,7 @@
 #include <stdio.h>
 #include <stdint.h>
 #include <stdlib.h>
+#include <string.h>
 #include <assert.h>
 #include <math.h>
 
@@ -8,12 +9,30 @@
 #include "ryg_rans/rans_byte.h"
 #include "ryg_rans/renormalize.h"
 
+#include <algorithm>
 #include <memory>
+#include <numeric>
+#include <random>
+#include <vector>
+#include <unordered_map>
 
 #define WIDTH 1280
 #define HEIGHT 720
+#define WIDTH_BLOCKS (WIDTH/8)
+#define WIDTH_BLOCKS_CHROMA (WIDTH/16)
+#define HEIGHT_BLOCKS (HEIGHT/8)
+#define NUM_BLOCKS (WIDTH_BLOCKS * HEIGHT_BLOCKS)
+#define NUM_BLOCKS_CHROMA (WIDTH_BLOCKS_CHROMA * HEIGHT_BLOCKS)
+
 #define NUM_SYMS 256
 #define ESCAPE_LIMIT (NUM_SYMS - 1)
+#define BLOCKS_PER_STREAM 320
+
+// If you set this to 1, the program will try to optimize the placement
+// of coefficients to rANS probability distributions. This is randomized,
+// so you might want to run it a few times.
+#define FIND_OPTIMAL_STREAM_ASSIGNMENT 0
+#define NUM_CLUSTERS 4
 
 static constexpr uint32_t prob_bits = 12;
 static constexpr uint32_t prob_scale = 1 << prob_bits;
@@ -67,7 +86,6 @@ struct SymbolStats
     uint32_t cum_freqs[NUM_SYMS + 1];
 
     void clear();
-    void count_freqs(uint8_t const* in, size_t nbytes);
     void calc_cum_freqs();
     void normalize_freqs(uint32_t target_total);
 };
@@ -78,14 +96,6 @@ void SymbolStats::clear()
         freqs[i] = 0;
 }
 
-void SymbolStats::count_freqs(uint8_t const* in, size_t nbytes)
-{
-    clear();
-
-    for (size_t i=0; i < nbytes; i++)
-        freqs[in[i]]++;
-}
-
 void SymbolStats::calc_cum_freqs()
 {
     cum_freqs[0] = 0;
@@ -93,119 +103,6 @@ void SymbolStats::calc_cum_freqs()
         cum_freqs[i+1] = cum_freqs[i] + freqs[i];
 }
 
-static double cache[NUM_SYMS + 1][prob_scale + 1];
-static double log2cache[prob_scale + 1];
-static int64_t cachefill = 0;
-
-double find_optimal_cost(const uint32_t *cum_freqs, int sym_to, int available_slots)
-{
-       assert(sym_to >= 0);
-
-       while (sym_to > 0 && cum_freqs[sym_to] == cum_freqs[sym_to - 1]) { --sym_to; }
-       if (cache[sym_to][available_slots] >= 0.0) {
-               //printf("CACHE: %d,%d\n", sym_to, available_slots);
-               return cache[sym_to][available_slots];
-       }
-       if (sym_to == 0) {
-               return 0.0;
-       }
-       if (sym_to == 1) {
-               return cum_freqs[0] * log2cache[available_slots];
-       }
-       if (available_slots == 1) {
-               return cum_freqs[0] * log2cache[1] + find_optimal_cost(cum_freqs, sym_to - 1, 0);
-       }
-
-//     printf("UNCACHE: %d,%d\n", sym_to, available_slots);
-#if 0
-       // ok, test all possible options for the last symbol (TODO: save the choice)
-       double best_so_far = HUGE_VAL;
-       //for (int i = num_syms - 1; i < available_slots; ++i) {
-       double f = freqs[sym_to - 1];
-       for (int i = available_slots; i --> 0; ) {
-               double cost1 = f * log2cache[available_slots - i];
-               double cost2 = find_optimal_cost(freqs, sym_to - 1, i);
-
-               if (sym_to == 3 && available_slots == 838) {
-                       printf("%d %f\n", i, cost1 + cost2);
-               } else
-               if (cost1 + cost2 > best_so_far) {
-                       break;
-               }
-               best_so_far = cost1 + cost2;
-       }
-#elif 1
-       // Minimize the number of total bits spent as a function of how many slots
-       // we assign to this symbol.
-       //
-       // The cost function is convex (I don't know how to prove it, but it makes
-       // intuitively a lot of sense). Find a reasonable guess and see what way
-       // we should search, then iterate until we either hit the end or we start
-       // increasing again.
-       double f = cum_freqs[sym_to - 1] - cum_freqs[sym_to - 2];
-       double start = lrint(available_slots * f / cum_freqs[sym_to - 1]);
-
-       int x1 = std::max<int>(floor(start), 1);
-       int x2 = x1 + 1;
-
-       double f1 = f * log2cache[x1] + find_optimal_cost(cum_freqs, sym_to - 1, available_slots - x1);
-       double f2 = f * log2cache[x2] + find_optimal_cost(cum_freqs, sym_to - 1, available_slots - x2);
-
-       int x, direction;  // -1 or +1
-       double best_so_far = std::min(f1, f2);
-       if (isinf(f1) && isinf(f2)) {
-               // The cost isn't infinite due to the first term, so we need to go downwards
-               // to give the second term more room to breathe.
-               x = x1;
-               direction = -1;
-       } else if (f1 < f2) {
-               x = x1;
-               direction = -1;
-       } else {
-               x = x2;
-               direction = 1;
-       }
-
-       //printf("[%d,%d] freq=%ld cumfreq=%d From %d and %d, chose %d [%f] and direction=%d\n",
-       //      sym_to, available_slots, freqs[sym_to - 1], cum_freqs[sym_to - 1], x1, x2, x, best_so_far, direction);
-
-       while ((x + direction) > 0 && (x + direction) <= available_slots) {
-               x += direction;
-               double fn = f * log2cache[x] + find_optimal_cost(cum_freqs, sym_to - 1, available_slots - x);
-       //      printf("[%d,%d] %d is %f\n", sym_to, available_slots, x, fn);
-               if (fn > best_so_far) {
-                       break;
-               }
-               best_so_far = fn;
-       }
-#endif
-       if (++cachefill % 131072 == 0) {
-       //      printf("%d,%d = %f (cachefill = %.2f%%)\n", sym_to, available_slots, best_so_far,
-       //              100.0 * (cachefill / double((NUM_SYMS + 1) * (prob_scale + 1))));
-       }
-       assert(best_so_far >= 0.0);
-       assert(cache[sym_to][available_slots] < 0.0);
-       cache[sym_to][available_slots] = best_so_far;
-       return best_so_far;
-}
-
-double find_optimal_cost(const uint32_t *cum_freqs, const uint64_t *freqs)
-{
-       for (int j = 0; j <= NUM_SYMS; ++j) {
-               for (unsigned k = 0; k <= prob_scale; ++k) {
-                       cache[j][k] = -1.0;
-               }
-       }
-       for (unsigned k = 0; k <= prob_scale; ++k) {
-               log2cache[k] = -log2(k * (1.0 / prob_scale));
-               //printf("log2cache[%d] = %f\n", k, log2cache[k]);
-       }
-       cachefill = 0;
-       double ret = find_optimal_cost(cum_freqs, NUM_SYMS, prob_scale);
-       printf("Used %ld function invocations\n", cachefill);
-       return ret;
-}
-
 void SymbolStats::normalize_freqs(uint32_t target_total)
 {
     uint64_t real_freq[NUM_SYMS + 1];  // hack
@@ -227,47 +124,6 @@ void SymbolStats::normalize_freqs(uint32_t target_total)
 
     OptimalRenormalize(cum_freqs, NUM_SYMS, prob_scale);
 
-#if 0
-    double optimal_cost = find_optimal_cost(cum_freqs + 1, real_freq + 1);
-
-    // resample distribution based on cumulative freqs
-    for (int i = 1; i <= NUM_SYMS; i++)
-        //cum_freqs[i] = ((uint64_t)target_total * cum_freqs[i])/cur_total;
-        cum_freqs[i] = lrint(cum_freqs[i] * double(target_total) / cur_total);
-
-    // if we nuked any non-0 frequency symbol to 0, we need to steal
-    // the range to make the frequency nonzero from elsewhere.
-    //
-    // this is not at all optimal, i'm just doing the first thing that comes to mind.
-    for (int i=0; i < NUM_SYMS; i++) {
-        if (freqs[i] && cum_freqs[i+1] == cum_freqs[i]) {
-            // symbol i was set to zero freq
-
-            // find best symbol to steal frequency from (try to steal from low-freq ones)
-            uint32_t best_freq = ~0u;
-            int best_steal = -1;
-            for (int j=0; j < NUM_SYMS; j++) {
-                uint32_t freq = cum_freqs[j+1] - cum_freqs[j];
-                if (freq > 1 && freq < best_freq) {
-                    best_freq = freq;
-                    best_steal = j;
-                }
-            }
-            assert(best_steal != -1);
-
-            // and steal from it!
-            if (best_steal < i) {
-                for (int j = best_steal + 1; j <= i; j++)
-                    cum_freqs[j]--;
-            } else {
-                assert(best_steal > i);
-                for (int j = i + 1; j <= best_steal; j++)
-                    cum_freqs[j]++;
-            }
-        }
-    }
-#endif
-
     // calculate updated freqs and make sure we didn't screw anything up
     assert(cum_freqs[0] == 0 && cum_freqs[NUM_SYMS] == target_total);
     for (int i=0; i < NUM_SYMS; i++) {
@@ -298,20 +154,43 @@ void SymbolStats::normalize_freqs(uint32_t target_total)
                 calc_cost, (calc_cost - ideal_cost) / 8.0, total_loss / 8.0, total_loss_with_dp / 8.0);
 }
 
-SymbolStats stats[64];
+SymbolStats stats[128];
 
-int pick_stats_for(int y, int x)
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+// Distance from one stream to the other, based on a hacked-up K-L divergence.
+float kl_dist[64][64];
+#endif
+
+const int luma_mapping[64] = {
+       0, 0, 1, 1, 2, 2, 3, 3,
+       0, 0, 1, 2, 2, 2, 3, 3,
+       1, 1, 2, 2, 2, 3, 3, 3,
+       1, 1, 2, 2, 2, 3, 3, 3,
+       1, 2, 2, 2, 2, 3, 3, 3,
+       2, 2, 2, 2, 3, 3, 3, 3,
+       2, 2, 3, 3, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3,
+};
+const int chroma_mapping[64] = {
+       0, 1, 1, 2, 2, 2, 3, 3,
+       1, 1, 2, 2, 2, 3, 3, 3,
+       2, 2, 2, 2, 3, 3, 3, 3,
+       2, 2, 2, 3, 3, 3, 3, 3,
+       2, 3, 3, 3, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3,
+};
+
+int pick_stats_for(int x, int y, bool is_chroma)
 {
-       //return std::min<int>(hypot(x, y), 7);
-       return std::min<int>(x + y, 7);
-       //if (x + y >= 7) return 7;
-       //return x + y;
-       //return y * 8 + x;
-#if 0
-       if (y == 0 && x == 0) {
-               return 0;
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+       return y * 8 + x + is_chroma * 64;
+#else
+       if (is_chroma) {
+               return chroma_mapping[y * 8 + x] + 4;
        } else {
-               return 1;
+               return luma_mapping[y * 8 + x];
        }
 #endif
 }
@@ -331,26 +210,23 @@ public:
        RansEncoder()
        {
                out_buf.reset(new uint8_t[out_max_size]);
-               sign_buf.reset(new uint8_t[max_num_sign]);
                clear();
        }
 
-       void init_prob(const SymbolStats &s)
+       void init_prob(SymbolStats &s)
        {
                for (int i = 0; i < NUM_SYMS; i++) {
-                       printf("%d: cumfreqs=%d freqs=%d prob_bits=%d\n", i, s.cum_freqs[i], s.freqs[i], prob_bits);
-                       RansEncSymbolInit(&esyms[i], s.cum_freqs[i], s.freqs[i], prob_bits);
+                       //printf("%d: cumfreqs=%d freqs=%d prob_bits=%d\n", i, s.cum_freqs[i], s.freqs[i], prob_bits + 1);
+                       RansEncSymbolInit(&esyms[i], s.cum_freqs[i], s.freqs[i], prob_bits + 1);
                }
+               sign_bias = s.cum_freqs[NUM_SYMS];
        }
 
        void clear()
        {
                out_end = out_buf.get() + out_max_size;
-               sign_end = sign_buf.get() + max_num_sign;
                ptr = out_end; // *end* of output buffer
-               sign_ptr = sign_end; // *end* of output buffer
                RansEncInit(&rans);
-               free_sign_bits = 0;
        }
 
        uint32_t save_block(FILE *codedfp)  // Returns number of bytes.
@@ -359,33 +235,33 @@ public:
                //printf("post-flush = %08x\n", rans);
 
                uint32_t num_rans_bytes = out_end - ptr;
+               if (num_rans_bytes == last_block.size() &&
+                   memcmp(last_block.data(), ptr, last_block.size()) == 0) {
+                       write_varint(0, codedfp);
+                       clear();
+                       return 1;
+               } else {
+                       last_block = string((const char *)ptr, num_rans_bytes);
+               }
+
                write_varint(num_rans_bytes, codedfp);
                //fwrite(&num_rans_bytes, 1, 4, codedfp);
                fwrite(ptr, 1, num_rans_bytes, codedfp);
 
                //printf("first rANS bytes: %02x %02x %02x %02x %02x %02x %02x %02x\n", ptr[0], ptr[1], ptr[2], ptr[3], ptr[4], ptr[5], ptr[6], ptr[7]);
 
-               if (free_sign_bits > 0) {
-                       *sign_ptr <<= free_sign_bits;
-               }
-
-#if 1
-               uint32_t num_sign_bytes = sign_end - sign_ptr;
-               write_varint((num_sign_bytes << 3) | free_sign_bits, codedfp);
-               fwrite(sign_ptr, 1, num_sign_bytes, codedfp);
-#endif
 
                clear();
 
-               printf("Saving block: %d rANS bytes, %d sign bytes\n", num_rans_bytes, num_sign_bytes);
-               return num_rans_bytes + num_sign_bytes;
+               //printf("Saving block: %d rANS bytes\n", num_rans_bytes);
+               return num_rans_bytes;
                //return num_rans_bytes;
        }
 
        void encode_coeff(short signed_k)
        {
-               printf("encoding coeff %d\n", signed_k);
-               short k = abs(signed_k);
+               //printf("encoding coeff %d (sym %d), rans before encoding = %08x\n", signed_k, ((abs(signed_k) - 1) & 255), rans);
+               unsigned short k = abs(signed_k);
                if (k >= ESCAPE_LIMIT) {
                        // Put the coefficient as a 1/(2^12) symbol _before_
                        // the 255 coefficient, since the decoder will read the
@@ -393,33 +269,24 @@ public:
                        RansEncPut(&rans, &ptr, k, 1, prob_bits);
                        k = ESCAPE_LIMIT;
                }
-               if (k != 0) {
-#if 1
-                       if (free_sign_bits == 0) {
-                               --sign_ptr;
-                               *sign_ptr = 0;
-                               free_sign_bits = 8;
-                       }
-                       *sign_ptr <<= 1;
-                       *sign_ptr |= (signed_k < 0);
-                       --free_sign_bits;
-#else
-                       RansEncPut(&rans, &ptr, (k < 0) ? prob_scale / 2 : 0, prob_scale / 2, prob_bits);
-#endif
+               RansEncPutSymbol(&rans, &ptr, &esyms[(k - 1) & (NUM_SYMS - 1)]);
+               if (signed_k < 0) {
+                       rans += sign_bias;
                }
-               RansEncPutSymbol(&rans, &ptr, &esyms[k]);
        }
 
 private:
        static constexpr size_t out_max_size = 32 << 20; // 32 MB.
        static constexpr size_t max_num_sign = 1048576;  // Way too big. And actually bytes.
 
-       unique_ptr<uint8_t[]> out_buf, sign_buf;
-       uint8_t *out_end, *sign_end;
-       uint8_t *ptr, *sign_ptr;
+       unique_ptr<uint8_t[]> out_buf;
+       uint8_t *out_end;
+       uint8_t *ptr;
        RansState rans;
-       size_t free_sign_bits;
        RansEncSymbol esyms[NUM_SYMS];
+       uint32_t sign_bias;
+
+       std::string last_block;
 };
 
 static constexpr int dc_scalefac = 8;  // Matches the FDCT's gain.
@@ -512,6 +379,114 @@ void convert_ycbcr()
        }
 }
 
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+double find_best_assignment(const int *medoids, int *assignment)
+{
+       double current_score = 0.0;
+       for (int i = 0; i < 64; ++i) {
+               int best_medoid = medoids[0];
+               float best_medoid_score = kl_dist[i][medoids[0]];
+               for (int j = 1; j < NUM_CLUSTERS; ++j) {
+                       if (kl_dist[i][medoids[j]] < best_medoid_score) {
+                               best_medoid = medoids[j];
+                               best_medoid_score = kl_dist[i][medoids[j]];
+                       }
+               }
+               assignment[i] = best_medoid;
+               current_score += best_medoid_score;
+       }
+       return current_score;
+}
+
+void find_optimal_stream_assignment(int base)
+{
+       double inv_sum[64];
+       for (unsigned i = 0; i < 64; ++i) {
+               double s = 0.0;
+               for (unsigned k = 0; k < NUM_SYMS; ++k) {
+                       s += stats[i + base].freqs[k] + 0.5;
+               }
+               inv_sum[i] = 1.0 / s;
+       }
+
+       for (unsigned i = 0; i < 64; ++i) {
+               for (unsigned j = 0; j < 64; ++j) {
+                       double d = 0.0;
+                       for (unsigned k = 0; k < NUM_SYMS; ++k) {
+                               double p1 = (stats[i + base].freqs[k] + 0.5) * inv_sum[i];
+                               double p2 = (stats[j + base].freqs[k] + 0.5) * inv_sum[j];
+
+                               // K-L divergence is asymmetric; this is a hack.
+                               d += p1 * log(p1 / p2);
+                               d += p2 * log(p2 / p1);
+                       }
+                       kl_dist[i][j] = d;
+                       //printf("%.3f ", d);
+               }
+               //printf("\n");
+       }
+
+       // k-medoids init
+       int medoids[64];  // only the first NUM_CLUSTERS matter
+       bool is_medoid[64] = { false };
+       std::iota(medoids, medoids + 64, 0);
+       std::random_device rd;
+       std::mt19937 g(rd());
+       std::shuffle(medoids, medoids + 64, g);
+       for (int i = 0; i < NUM_CLUSTERS; ++i) {
+               printf("%d ", medoids[i]);
+               is_medoid[medoids[i]] = true;
+       }
+       printf("\n");
+
+       // assign each data point to the closest medoid
+       int assignment[64];
+       double current_score = find_best_assignment(medoids, assignment);
+
+       for (int i = 0; i < 1000; ++i) {
+               printf("iter %d\n", i);
+               bool any_changed = false;
+               for (int m = 0; m < NUM_CLUSTERS; ++m) {
+                       for (int o = 0; o < 64; ++o) {
+                               if (is_medoid[o]) continue;
+                               int old_medoid = medoids[m];
+                               medoids[m] = o;
+
+                               int new_assignment[64];
+                               double candidate_score = find_best_assignment(medoids, new_assignment);
+
+                               if (candidate_score < current_score) {
+                                       current_score = candidate_score;
+                                       memcpy(assignment, new_assignment, sizeof(assignment));
+
+                                       is_medoid[old_medoid] = false;
+                                       is_medoid[medoids[m]] = true;
+                                       printf("%f: ", current_score);
+                                       for (int i = 0; i < 64; ++i) {
+                                               printf("%d ", assignment[i]);
+                                       }
+                                       printf("\n");
+                                       any_changed = true;
+                               } else {
+                                       medoids[m] = old_medoid;
+                               }
+                       }
+               }
+               if (!any_changed) break;
+       }
+       printf("\n");
+       std::unordered_map<int, int> rmap;
+       for (int i = 0; i < 64; ++i) {
+               if (i % 8 == 0) printf("\n");
+               if (!rmap.count(assignment[i])) {
+                       rmap.emplace(assignment[i], rmap.size());
+               }
+               printf("%d, ", rmap[assignment[i]]);
+       }
+       printf("\n");
+}
+#endif
+
 int main(int argc, char **argv)
 {
        if (argc >= 2)
@@ -524,6 +499,9 @@ int main(int argc, char **argv)
        //double last_cb_cfl_fac = 0.0;
        //double last_cr_cfl_fac = 0.0;
 
+       int max_val_x[8] = {0}, min_val_x[8] = {0};
+       int max_val_y[8] = {0}, min_val_y[8] = {0};
+
        // DCT and quantize luma
        for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
                for (unsigned xb = 0; xb < WIDTH; xb += 8) {
@@ -544,6 +522,11 @@ int main(int argc, char **argv)
                                        int k = quantize(in_y[coeff_idx], coeff_idx);
                                        coeff_y[(yb + y) * WIDTH + (xb + x)] = k;
 
+                                       max_val_x[x] = std::max(max_val_x[x], k);
+                                       min_val_x[x] = std::min(min_val_x[x], k);
+                                       max_val_y[y] = std::max(max_val_y[y], k);
+                                       min_val_y[y] = std::min(min_val_y[y], k);
+
                                        // Store back for reconstruction / PSNR calculation
                                        in_y[coeff_idx] = unquantize(k, coeff_idx);
                                }
@@ -738,10 +721,35 @@ int main(int argc, char **argv)
                chroma_energy / (WIDTH * HEIGHT), chroma_energy_pred / (WIDTH * HEIGHT));
 #endif
 
-       // DC coefficient pred from the right to left
-       for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
-               for (unsigned xb = 0; xb < WIDTH - 8; xb += 8) {
-                       coeff_y[yb * WIDTH + xb] -= coeff_y[yb * WIDTH + (xb + 8)];
+       // DC coefficient pred from the right to left (within each slice)
+       for (unsigned block_idx = 0; block_idx < NUM_BLOCKS; block_idx += BLOCKS_PER_STREAM) {
+               int prev_k = 128;
+
+               for (unsigned subblock_idx = BLOCKS_PER_STREAM; subblock_idx --> 0; ) {
+                       unsigned yb = (block_idx + subblock_idx) / WIDTH_BLOCKS;
+                       unsigned xb = (block_idx + subblock_idx) % WIDTH_BLOCKS;
+                       int k = coeff_y[(yb * 8) * WIDTH + (xb * 8)];
+
+                       coeff_y[(yb * 8) * WIDTH + (xb * 8)] = k - prev_k;
+
+                       prev_k = k;
+               }
+       }
+       for (unsigned block_idx = 0; block_idx < NUM_BLOCKS_CHROMA; block_idx += BLOCKS_PER_STREAM) {
+               int prev_k_cb = 0;
+               int prev_k_cr = 0;
+
+               for (unsigned subblock_idx = BLOCKS_PER_STREAM; subblock_idx --> 0; ) {
+                       unsigned yb = (block_idx + subblock_idx) / WIDTH_BLOCKS_CHROMA;
+                       unsigned xb = (block_idx + subblock_idx) % WIDTH_BLOCKS_CHROMA;
+                       int k_cb = coeff_cb[(yb * 8) * WIDTH/2 + (xb * 8)];
+                       int k_cr = coeff_cr[(yb * 8) * WIDTH/2 + (xb * 8)];
+
+                       coeff_cb[(yb * 8) * WIDTH/2 + (xb * 8)] = k_cb - prev_k_cb;
+                       coeff_cr[(yb * 8) * WIDTH/2 + (xb * 8)] = k_cr - prev_k_cr;
+
+                       prev_k_cb = k_cb;
+                       prev_k_cr = k_cr;
                }
        }
 
@@ -778,33 +786,31 @@ int main(int argc, char **argv)
        fclose(fp);
 
        // For each coefficient, make some tables.
-       size_t extra_bits = 0, sign_bits = 0;
+       size_t extra_bits = 0;
        for (unsigned i = 0; i < 64; ++i) {
                stats[i].clear();
        }
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_luma = stats[pick_stats_for(x, y)];
-                       SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];  // HACK
-                       //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_luma = stats[pick_stats_for(x, y, false)];
+                       SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
 
                        // Luma
                        for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
                                for (unsigned xb = 0; xb < WIDTH; xb += 8) {
-                                       short k = abs(coeff_y[(yb + y) * WIDTH + (xb + x)]);
+                                       unsigned short k = abs(coeff_y[(yb + y) * WIDTH + (xb + x)]);
                                        if (k >= ESCAPE_LIMIT) {
                                                k = ESCAPE_LIMIT;
                                                extra_bits += 12;  // escape this one
                                        }
-                                       if (k != 0) ++sign_bits;
-                                       ++s_luma.freqs[k];
+                                       ++s_luma.freqs[(k - 1) & (NUM_SYMS - 1)];
                                }
                        }
                        // Chroma
                        for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
                                for (unsigned xb = 0; xb < WIDTH/2; xb += 8) {
-                                       short k_cb = abs(coeff_cb[(yb + y) * WIDTH/2 + (xb + x)]);
-                                       short k_cr = abs(coeff_cr[(yb + y) * WIDTH/2 + (xb + x)]);
+                                       unsigned short k_cb = abs(coeff_cb[(yb + y) * WIDTH/2 + (xb + x)]);
+                                       unsigned short k_cr = abs(coeff_cr[(yb + y) * WIDTH/2 + (xb + x)]);
                                        if (k_cb >= ESCAPE_LIMIT) {
                                                k_cb = ESCAPE_LIMIT;
                                                extra_bits += 12;  // escape this one
@@ -813,16 +819,26 @@ int main(int argc, char **argv)
                                                k_cr = ESCAPE_LIMIT;
                                                extra_bits += 12;  // escape this one
                                        }
-                                       if (k_cb != 0) ++sign_bits;
-                                       if (k_cr != 0) ++sign_bits;
-                                       ++s_chroma.freqs[k_cb];
-                                       ++s_chroma.freqs[k_cr];
+                                       ++s_chroma.freqs[(k_cb - 1) & (NUM_SYMS - 1)];
+                                       ++s_chroma.freqs[(k_cr - 1) & (NUM_SYMS - 1)];
                                }
                        }
                }
        }
+
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+       printf("Luma:\n");
+       find_optimal_stream_assignment(0);
+       printf("Chroma:\n");
+       find_optimal_stream_assignment(64);
+       exit(0);
+#endif
+
        for (unsigned i = 0; i < 64; ++i) {
+               stats[i].freqs[NUM_SYMS - 1] /= 2;  // zero, has no sign bits (yes, this is trickery)
                stats[i].normalize_freqs(prob_scale);
+               stats[i].cum_freqs[NUM_SYMS] += stats[i].freqs[NUM_SYMS - 1];
+               stats[i].freqs[NUM_SYMS - 1] *= 2;
        }
 
        FILE *codedfp = fopen("coded.dat", "wb");
@@ -849,53 +865,67 @@ int main(int argc, char **argv)
        // Luma
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_luma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_luma = stats[pick_stats_for(x, y, false)];
                        rans_encoder.init_prob(s_luma);
 
                        // Luma
+                       std::vector<int> lens;
 
                        // need to reverse later
                        rans_encoder.clear();
                        size_t num_bytes = 0;
-                       for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
-                               for (unsigned xb = 0; xb < WIDTH; xb += 8) {
-                                       int k = coeff_y[(yb + y) * WIDTH + (xb + x)];
-                                       //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
-                                       rans_encoder.encode_coeff(k);
-                               }
-                               if (yb % 16 == 8) {
-                                       num_bytes += rans_encoder.save_block(codedfp);
+                       for (unsigned block_idx = 0; block_idx < NUM_BLOCKS; ++block_idx) {
+                               unsigned yb = block_idx / WIDTH_BLOCKS;
+                               unsigned xb = block_idx % WIDTH_BLOCKS;
+
+                               int k = coeff_y[(yb * 8 + y) * WIDTH + (xb * 8 + x)];
+                               //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
+                               rans_encoder.encode_coeff(k);
+
+                               if (block_idx % BLOCKS_PER_STREAM == (BLOCKS_PER_STREAM - 1) || block_idx == NUM_BLOCKS - 1) {
+                                       int l = rans_encoder.save_block(codedfp);
+                                       num_bytes += l;
+                                       lens.push_back(l);
                                }
                        }
-                       if (HEIGHT % 16 != 0) {
-                               num_bytes += rans_encoder.save_block(codedfp);
-                       }
                        tot_bytes += num_bytes;
                        printf("coeff %d Y': %ld bytes\n", y * 8 + x, num_bytes);
+
+                       double sum_l = 0.0;
+                       for (int l : lens) {
+                               sum_l += l;
+                       }
+                       double avg_l = sum_l / lens.size();
+
+                       double sum_sql = 0.0;
+                       for (int l : lens) {
+                               sum_sql += (l - avg_l) * (l - avg_l);
+                       }
+                       double stddev_l = sqrt(sum_sql / (lens.size() - 1));
+                       printf("coeff %d: avg=%.2f bytes, stddev=%.2f bytes\n", y*8+x, avg_l, stddev_l);
                }
        }
 
        // Cb
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];
-                       //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
                        rans_encoder.init_prob(s_chroma);
 
                        rans_encoder.clear();
                        size_t num_bytes = 0;
-                       for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
-                               for (unsigned xb = 0; xb < WIDTH/2; xb += 8) {
-                                       int k = coeff_cb[(yb + y) * WIDTH/2 + (xb + x)];
-                                       rans_encoder.encode_coeff(k);
-                               }
-                               if (yb % 16 == 8) {
+                       for (unsigned block_idx = 0; block_idx < NUM_BLOCKS_CHROMA; ++block_idx) {
+                               unsigned yb = block_idx / WIDTH_BLOCKS_CHROMA;
+                               unsigned xb = block_idx % WIDTH_BLOCKS_CHROMA;
+
+                               int k = coeff_cb[(yb * 8 + y) * WIDTH/2 + (xb * 8 + x)];
+                               //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
+                               rans_encoder.encode_coeff(k);
+
+                               if (block_idx % BLOCKS_PER_STREAM == (BLOCKS_PER_STREAM - 1) || block_idx == NUM_BLOCKS - 1) {
                                        num_bytes += rans_encoder.save_block(codedfp);
                                }
                        }
-                       if (HEIGHT % 16 != 0) {
-                               num_bytes += rans_encoder.save_block(codedfp);
-                       }
                        tot_bytes += num_bytes;
                        printf("coeff %d Cb: %ld bytes\n", y * 8 + x, num_bytes);
                }
@@ -904,34 +934,45 @@ int main(int argc, char **argv)
        // Cr
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];
-                       //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
                        rans_encoder.init_prob(s_chroma);
 
                        rans_encoder.clear();
                        size_t num_bytes = 0;
-                       for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
-                               for (unsigned xb = 0; xb < WIDTH/2; xb += 8) {
-                                       int k = coeff_cr[(yb + y) * WIDTH/2 + (xb + x)];
-                                       rans_encoder.encode_coeff(k);
-                               }
-                               if (yb % 16 == 8) {
+                       for (unsigned block_idx = 0; block_idx < NUM_BLOCKS_CHROMA; ++block_idx) {
+                               unsigned yb = block_idx / WIDTH_BLOCKS_CHROMA;
+                               unsigned xb = block_idx % WIDTH_BLOCKS_CHROMA;
+
+                               int k = coeff_cr[(yb * 8 + y) * WIDTH/2 + (xb * 8 + x)];
+                               //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
+                               rans_encoder.encode_coeff(k);
+
+                               if (block_idx % BLOCKS_PER_STREAM == (BLOCKS_PER_STREAM - 1) || block_idx == NUM_BLOCKS - 1) {
                                        num_bytes += rans_encoder.save_block(codedfp);
                                }
                        }
-                       if (HEIGHT % 16 != 0) {
-                               num_bytes += rans_encoder.save_block(codedfp);
-                       }
                        tot_bytes += num_bytes;
                        printf("coeff %d Cr: %ld bytes\n", y * 8 + x, num_bytes);
                }
        }
 
-       printf("%ld bytes + %ld sign bits (%ld) + %ld escape bits (%ld) = %ld total bytes\n",
-               tot_bytes - sign_bits / 8 - extra_bits / 8,
-               sign_bits,
-               sign_bits / 8,
+       printf("%ld bytes + %ld escape bits (%ld) = %ld total bytes\n",
+               tot_bytes - extra_bits / 8,
                extra_bits,
                extra_bits / 8,
                tot_bytes);
+
+#if 0
+       printf("Max coefficient ranges (as a function of x):\n\n");
+       for (unsigned x = 0; x < 8; ++x) {
+               int range = std::max(max_val_x[x], -min_val_x[x]);
+               printf("  [%4d, %4d] (%.2f bits)\n", min_val_x[x], max_val_x[x], log2(range * 2 + 1));
+       }
+
+       printf("Max coefficient ranges (as a function of y):\n\n");
+       for (unsigned y = 0; y < 8; ++y) {
+               int range = std::max(max_val_y[y], -min_val_y[y]);
+               printf("  [%4d, %4d] (%.2f bits)\n", min_val_y[y], max_val_y[y], log2(range * 2 + 1));
+       }
+#endif
 }