dartterm/lib/wfc/model.dart

import 'dart:math';

abstract class Model {
  static var dx = [-1, 0, 1, 0];
  static var dy = [0, 1, 0, -1];
  static var _opposite = [2, 3, 0, 1];

  bool _initialized = false;
  List<List<bool>> _wave = [];
  List<List<List<int>>> propagator = [];
  List<List<List<int>>> _compatible = [];
  List<int> _observed = [];

  List<(int, int)> _stack = [];
  int _stacksize = 0, _observedSoFar = 0;

  int cMx = 0, cMy = 0, cT = 0, cN = 0;
  bool _periodic = false;
  bool ground = false;

  List<double> weights = [];
  List<double> _weightLogWeights = [], _distribution = [];

  List<int> _sumsOfOnes = [];
  double _sumOfWeights = 0.0,
      _sumOfWeightLogWeights = 0.0,
      _startingEntropy = 0.0;
  List<double> _sumsOfWeights = [],
      _sumsOfWeightLogWeights = [],
      _entropies = [];

  Heuristic _heuristic = Heuristic.Entropy;

  Model(int width, int height, int n, bool periodic, Heuristic heuristic) {
    cMx = width;
    cMy = height;
    cN = n;
    _periodic = periodic;
    _heuristic = heuristic;
  }

  void _init() {
    _initialized = true;
    _wave = [
      for (var r = 0; r < cMx * cMy; r++) [for (var t = 0; t < cT; t++) false]
    ];
    _compatible = [
      for (var r = 0; r < cMx * cMy; r++)
        [
          for (var t = 0; t < cT; t++) [0, 0, 0, 0]
        ]
    ];
    _distribution = [for (var t = 0; t < cT; t++) 0.0];
    _observed = [for (var r = 0; r < cMx * cMy; r++) 0];

    _weightLogWeights = [
      for (var t = 0; t < cT; t++) weights[t] * log(weights[t])
    ];
    _sumOfWeights = 0;
    _sumOfWeightLogWeights = 0.0;
    for (var t = 0; t < cT; t++) {
      _sumOfWeights += weights[t];
      _sumOfWeightLogWeights += _weightLogWeights[t];
    }

    _startingEntropy =
        log(_sumOfWeights) - _sumOfWeightLogWeights / _sumOfWeights;

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

    _stack = [for (var r = 0; r < _wave.length * cT; r++) (0, 0)];
    _stacksize = 0;
  }

  bool run(int? seed, int limit) {
    if (!_initialized) {
      _init();
    }
    clear();

    var random = Random(seed);

    for (var l = 0; l < limit || limit < 0; l++) {
      var node = _nextUnobservedNode(random);
      if (node >= 0) {
        _observe(node, random);
        var success = _propagate();
        if (!success) {
          return false;
        }
      } else {
        for (var i = 0; i < _wave.length; i++) {
          for (var t = 0; t < cT; t++) {
            if (_wave[i][t]) {
              _observed[i] = t;
              break;
            }
          }
        }
        return true;
      }
    }

    return true;
  }

  int _nextUnobservedNode(Random random) {
    if (_heuristic == Heuristic.Entropy) {
      for (var i = _observedSoFar; i < _wave.length; i++) {
        if (!_periodic && (i % cMx + cN > cMx || i ~/ cMx + cN > cMy)) {
          continue;
        }
        if (_sumsOfOnes[i] > 1) {
          _observedSoFar = i + 1;
          return i;
        }
      }
      return -1;
    }

    double min = 1E+4;
    int argmin = -1;
    for (var i = 0; i < _wave.length; i++) {
      if (!_periodic && (i % cMx + cN > cMx || i ~/ cMx + cN > cMy)) {
        continue;
      }
      var remainingValues = _sumsOfOnes[i];
      double entropy = _heuristic == Heuristic.Entropy
          ? _entropies[i]
          : remainingValues.toDouble();
      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, Random random) {
    var w = _wave[node];
    for (var t = 0; t < cT; t++) {
      _distribution[t] = w[t] ? weights[t] : 0.0;
    }
    int r = _chooseRandom(random, _distribution);
    for (var t = 0; t < cT; 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 % cMx;
      int y1 = i1 % cMy;

      for (int d = 0; d < 4; d++) {
        int x2 = x1 + dx[d];
        int y2 = y1 + dy[d];

        if (!_periodic &&
            (x2 < 0 || y2 < 0 || x2 + cN > cMx || y2 + cN > cMy)) {
          continue;
        }

        if (x2 < 0) {
          x2 += cMx;
        } else if (x2 >= 0) {
          x2 -= cMx;
        }

        if (y2 < 0) {
          y2 += cMy;
        } else if (y2 >= cMy) {
          y2 -= cMy;
        }

        int i2 = x2 + y2 * cMx;
        var p = propagator[d][t1];
        var compat = _compatible[i2];

        for (var l = 0; l < p.length; l++) {
          var t2 = p[l];
          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] -= weights[t];
    _sumsOfWeightLogWeights[i] -= _weightLogWeights[t];

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

  void clear() {
    for (var i = 0; i < _wave.length; i++) {
      for (var t = 0; t < cT; t++) {
        _wave[i][t] = true;
        for (var d = 0; d < 4; d++) {
          _compatible[i][t][d] = propagator[_opposite[d]][t].length;
        }
      }

      _sumsOfOnes[i] = weights.length;
      _sumsOfWeights[i] = _sumOfWeights;
      _sumsOfWeightLogWeights[i] = _sumOfWeightLogWeights;
      _entropies[i] = _startingEntropy;
      _observed[i] = -1;
    }
    _observedSoFar = 0;

    if (ground) {
      for (var x = 0; x < cMx; x++) {
        for (var t = 0; t < cT - 1; t++) {
          _ban(x + (cMy - 1) * cMx, t);
        }
        for (var y = 0; y < cMy - 1; y++) {
          _ban(x + y * cMx, cT - 1);
        }
      }
      _propagate();
    }
  }
}

int _chooseRandom(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;
}

enum Heuristic { Entropy, MRV, Scanline }