dartterm/lib/wfc/template.dart

import 'dart:developer';
import 'dart:math' as math;
import 'dart:typed_data';
import 'dart:ui' as ui;
import 'dart:ui';

import 'package:flutter/services.dart';

class Wfc<T> {
  // parameters
  final WfcTemplate<T> _template;
  final int _mx, _my;

  // constants
  int get _n => _mx * _my;
  int get _nShingles => _template._shingles.n;
  int get _order => _template._order;
  double _weight(int shingleIx) =>
      _template._shingles._shingleWeights[shingleIx];

  // overall algo state
  List<List<bool>> _wave = [];
  List<List<List<int>>> _compatible = [];
  List<int?> _observed = [];

  // computationally expensive stuff that we keep in an incremental way
  List<double> _weightLogWeights = [];
  double _sumOfWeights = 0.0, _sumOfWeightLogWeights = 0.0;
  double _startingEntropy = 0.0;

  List<int> _sumsOfOnes = [];
  List<double> _sumsOfWeights = [];
  List<double> _sumsOfWeightLogWeights = [];
  List<double> _entropies = [];

  // temporaries
  List<double> _distribution = [];
  List<(int, int)> _stack = [];
  int _stacksize = 0;

  Wfc(this._template, this._mx, this._my) {
    _wave = [
      for (var r = 0; r < _n; r++) [for (var t = 0; t < _nShingles; t++) false]
    ];

    _compatible = [
      for (var r = 0; r < _n; r++)
        [
          for (var t = 0; t < _nShingles; t++) [0, 0, 0, 0]
        ]
    ];

    _distribution = [for (var t = 0; t < _nShingles; t++) 0.0];
    _observed = [for (var r = 0; r < _n; r++) null];

    _weightLogWeights = [
      for (var t = 0; t < _nShingles; t++) _weight(t) * math.log(_weight(t))
    ];
    _sumOfWeights = 0.0;
    _sumOfWeightLogWeights = 0.0;
    for (var t = 0; t < _nShingles; t++) {
      _sumOfWeights += _weight(t);
      _sumOfWeightLogWeights += _weightLogWeights[t];
    }
    _startingEntropy =
        math.log(_sumOfWeights) - _sumOfWeightLogWeights / _sumOfWeights;

    _sumsOfOnes = [for (var r = 0; r < _n; r++) 0];
    _sumsOfWeights = [for (var r = 0; r < _n; r++) 0.0];
    _sumsOfWeightLogWeights = [for (var r = 0; r < _n; r++) 0.0];
    _entropies = [for (var r = 0; r < _n; r++) 0.0];

    _stack = [for (var r = 0; r < _n * _nShingles; r++) (0, 0)];
    _stacksize = 0;
  }

  void clear() {
    for (var i = 0; i < _wave.length; i++) {
      for (var t = 0; t < _nShingles; t++) {
        _wave[i][t] = true;
        for (var d = 0; d < 4; d++) {
          _compatible[i][t][d] = _template
              ._shingles._metadata._propagators[_opposite[d]][t].length;
        }
      }

      _sumsOfOnes[i] = _nShingles;
      _sumsOfWeights[i] = _sumOfWeights;
      _sumsOfWeightLogWeights[i] = _sumOfWeightLogWeights;
      _entropies[i] = _startingEntropy;
      _observed[i] = null;
    }
  }

  bool run(int? seed, int limit) {
    clear();

    var random = math.Random(seed);
    for (var l = 0; l < limit || limit < 0; l++) {
      var node = _nextUnobservedNode(random);
      if (node != null) {
        _observe(node, random);
        var success = _propagate();
        if (!success) {
          return false;
        }
      } else {
        for (var i = 0; i < _n; i++) {
          for (var t = 0; t < _nShingles; t++) {
            if (_wave[i][t]) {
              _observed[i] = t;
              break;
            }
          }
        }
        return true;
      }
    }

    return true;
  }

  List<T>? extract() {
    var partial = extractPartial();
    List<T> out = [];
    for (var i in partial) {
      if (i == null) {
        return null;
      }
      out.add(i);
    }
    return out;
  }

  List<T?> extractPartial() {
    List<T?> result = [];
    for (int i = 0; i < _n; i++) {
      result.add(null);
    }

    for (int y = 0; y < _my; y++) {
      var dy = y < _my - _order + 1 ? 0 : _order - 1;
      for (int x = 0; x < _mx; x++) {
        var dx = x < _mx - _order + 1 ? 0 : _order - 1;
        var shingleIx = _observed[x - dx + (y - dy) * _mx];
        if (shingleIx != null) {
          var shingle = _template._shingles._shingleValues[shingleIx];
          var content = shingle.content[dx + dy * _order];
          var real = _template._embedding.decode(content);

          result[x + y * _mx] = real;
        }
      }
    }

    return result;
  }

  int? _nextUnobservedNode(math.Random random) {
    double min = 1E+10;
    int? argmin;
    for (var i = 0; i < _n; i++) {
      if (i % _mx + _order > _mx || i ~/ _mx + _order > _my) {
        continue;
      }
      var remainingValues = _sumsOfOnes[i];
      double entropy = remainingValues.toDouble(); // _entropies[i];
      if (remainingValues > 1 && entropy <= min) {
        double noise = 1E-6 * random.nextDouble();
        if (entropy + noise < min) {
          min = entropy + noise;
          argmin = i;
        }
      }
    }
    return argmin;
  }

  void _observe(int node, math.Random random) {
    var w = _wave[node];
    for (var t = 0; t < _nShingles; t++) {
      _distribution[t] = w[t] ? _weight(t) : 0.0;
    }

    int r = _chooseRandom(random, _distribution);
    for (var t = 0; t < _nShingles; t++) {
      if (w[t] != (t == r)) {
        _ban(node, t);
      }
    }
  }

  bool _propagate() {
    while (_stacksize > 0) {
      int i1, t1;
      (i1, t1) = _stack[_stacksize - 1];
      _stacksize--;

      int x1 = i1 % _mx;
      int y1 = i1 ~/ _mx;

      for (int d = 0; d < 4; d++) {
        var x2 = x1 + _dx[d];
        var y2 = y1 + _dy[d];

        if (x2 < 0 || y2 < 0 || x2 + _order > _mx || y2 + _order > _my) {
          continue;
        }

        int i2 = x2 + y2 * _mx;
        var p = _template._shingles._metadata._propagators[d][t1];
        var compat = _compatible[i2];

        for (var t2 in p) {
          var comp = compat[t2];
          comp[d]--;
          if (comp[d] == 0) {
            _ban(i2, t2);
          }
        }
      }
    }

    return _sumsOfOnes[0] > 0;
  }

  void _ban(int i, int t) {
    _wave[i][t] = false;

    var comp = _compatible[i][t];
    for (var d = 0; d < 4; d++) {
      comp[d] = 0;
    }
    _stack[_stacksize] = (i, t);
    _stacksize++;

    _sumsOfOnes[i] -= 1;
    _sumsOfWeights[i] -= _weight(t);
    _sumsOfWeightLogWeights[i] -= _weightLogWeights[t];

    var sum = _sumsOfWeights[i];
    _entropies[i] = math.log(sum) - _sumsOfWeightLogWeights[i] / sum;
  }
}

class WfcTemplate<T> {
  final Shingles _shingles;
  final Embedding<T> _embedding;
  final int _order;

  WfcTemplate(this._shingles, this._embedding, this._order);

  static Future<WfcTemplate<T>> loadAsync<T>(
      String name, int order, T Function(int) cb) async {
    final assetImageByteData = await rootBundle.load(name);
    final codec =
        await ui.instantiateImageCodec(assetImageByteData.buffer.asUint8List());
    final image = (await codec.getNextFrame()).image;

    final bytedata =
        (await image.toByteData(format: ImageByteFormat.rawStraightRgba))!;

    final sx = image.width;
    final sy = image.height;

    final List<T> bitmap = [];
    for (var i = 0; i < sx * sy; i++) {
      var pixel = bytedata.getUint32(i * 4, Endian.little);
      log("pixel: $pixel");
      bitmap.add(cb(pixel));
    }

    return loadBitmap(bitmap, sx, sy, order);
  }

  static WfcTemplate<T> loadBitmap<T>(
      List<T> bitmap, int sx, int sy, int order) {
    if (bitmap.length != sx * sy) {
      throw Exception("malformed bitmap");
    }
    var embedding = Embedding<T>();
    List<int> sample = [
      for (var i = 0; i < bitmap.length; i++) embedding.encode(bitmap[i])
    ];
    embedding.freeze();

    var shingles = Shingles();
    var xmax = sx - order + 1;
    var ymax = sy - order + 1;
    for (var y = 0; y < ymax; y++) {
      for (var x = 0; x < xmax; x++) {
        var ps = [
          Shingle(order, embedding.c,
              (dx, dy) => sample[(x + dx) % sx + (y + dy) % sy * sx])
        ];
        ps.add(ps[0].reflect());
        ps.add(ps[0].rotate());
        ps.add(ps[2].reflect());
        ps.add(ps[2].rotate());
        ps.add(ps[4].reflect());
        ps.add(ps[4].rotate());
        ps.add(ps[6].reflect());
        for (var p in ps) {
          shingles.observe(p, 1.0);
        }
      }
    }
    shingles.freeze();

    return WfcTemplate(shingles, embedding, order);
  }
}

class Shingles {
  bool _frozen = false;
  final Map<int, int> _shingleIndices = {};
  final List<Shingle> _shingleValues = [];
  final List<double> _shingleWeights = [];

  late ShingleMetadata _metadata;

  int get n {
    if (!_frozen) {
      throw StateError("can't use Shingles#get n until frozen");
    }
    return _shingleValues.length;
  }

  void freeze() {
    if (_frozen) {
      throw StateError("can't freeze when already frozen");
    }
    _frozen = true;

    _metadata = ShingleMetadata(this);
  }

  void observe(Shingle s, double n) {
    if (_frozen) {
      throw StateError("can't observe when already frozen");
    }

    // double n: weights can be fractional
    var index = _shingleIndices[s.hashCode];
    if (index == null) {
      index = _shingleValues.length;
      _shingleValues.add(s);
      _shingleWeights.add(n);
    } else {
      _shingleWeights[index] += n;
    }
  }
}

class ShingleMetadata {
  // [direction][source] => list of agreeing items
  List<List<List<int>>> _propagators = [];

  ShingleMetadata(Shingles s) {
    _propagators = [
      for (var d = 0; d < 4; d++)
        [
          for (var t = 0; t < s.n; t++)
            [
              for (var t2 = 0; t2 < s.n; t2++)
                if (s._shingleValues[t]
                    .agrees(s._shingleValues[t2], _dx[d], _dy[d]))
                  t2
            ]
        ]
    ];
  }
}

class Shingle {
  int order;
  int c;
  List<int> content = [];

  @override
  int hashCode = 0;

  Shingle(this.order, this.c, int Function(int, int) f) {
    content = [
      for (var y = 0; y < order; y++)
        for (var x = 0; x < order; x++) f(x, y)
    ];

    int result = 0, power = 1;
    for (var i = 0; i < content.length; i++) {
      result += content[content.length - 1 - i] * power;
      power *= c;
    }
    hashCode = result;
  }

  Shingle rotate() {
    return Shingle(order, c, (x, y) => content[order - 1 - y + x * order]);
  }

  Shingle reflect() {
    return Shingle(order, c, (x, y) => content[order - 1 - x + y * order]);
  }

  bool agrees(Shingle other, int dx, int dy) {
    var p1 = content;
    var p2 = other.content;
    var n = order;

    int xmin = dx < 0 ? 0 : dx;
    int xmax = dx < 0 ? dx + n : n;
    int ymin = dy < 0 ? 0 : dy;
    int ymax = dy < 0 ? dy + n : n;
    for (var y = ymin; y < ymax; y++) {
      for (var x = xmin; x < xmax; x++) {
        if (p1[x + n * y] != p2[x - dx + n * (y - dy)]) {
          return false;
        }
      }
    }
    return true;
  }

  @override
  bool operator ==(Object other) {
    return (other is Shingle) &&
        other.hashCode == hashCode &&
        other.order == order &&
        other.c == c;
  }
}

class Embedding<T> {
  bool _frozen = false;
  final List<T> _colorOf = [];
  final Map<T, int> _codeOf = {};

  void freeze() {
    if (_frozen) {
      throw StateError("can't freeze when already frozen");
    }
    _frozen = true;
  }

  int get c {
    if (!_frozen) {
      throw StateError("can't use Embedding#get c until frozen");
    }
    return _colorOf.length;
  }

  int encode(T t) {
    var code = _codeOf[t];
    if (code == null) {
      if (_frozen) {
        throw StateError("can't create new code when frozen");
      }
      code = _colorOf.length;
      _codeOf[t] = code;
      _colorOf.add(t);
    }
    return code;
  }

  T decode(int i) {
    return _colorOf[i]!;
  }
}

int _chooseRandom(math.Random rand, List<double> distribution) {
  if (distribution.isEmpty) {
    throw Exception("can't sample empty distribution");
  }
  var sum = 0.0;
  for (var i = 0; i < distribution.length; i++) {
    sum += distribution[i];
  }

  if (sum == 0.0) {
    return rand.nextInt(distribution.length);
  }

  var rnd = rand.nextDouble() * sum;

  var i = 0;
  while (rnd > 0) {
    rnd -= distribution[i];
    if (rnd < 0) {
      return i;
    }
    i += 1;
  }
  return distribution.length - 1;
}

final List<int> _dx = [-1, 0, 1, 0];
final List<int> _dy = [0, 1, 0, -1];
final List<int> _opposite = [2, 3, 0, 1];