fflame.nt

module fflame;

macro import std.macro.assert;
macro import std.macro.cimport;
macro import std.macro.listcomprehension;

import c_header("SDL2/SDL.h");
import c_header("SDL2/SDL_image.h");
import std.math;
import std.math.vector;
import std.string;
import std.thread;

pragma(lib, "SDL2");
pragma(lib, "SDL2_image");

class XorShift
{
    mut ulong a;
    this(ulong seed) { a = seed; }
    ulong rand() {
        mut ulong x = a;
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
        a = x;
        return x;
    }
    int irand() {
        return cast(int) (rand % 1073741824);
    }
    float randf() {
        return irand / 1073741824.0f;
    }
}

extern(C) void* neat_runtime_stdout();
extern(C) void print(string);
extern(C) char* toStringz(string);
extern(C) long atol(char* ptr);

struct Point
{
    float x, y;
    void weightAdd(Point p2, float weight) {
        this.x += p2.x * weight;
        this.y += p2.y * weight;
    }
}

AffineTransform makeAffineTransform(float a, float b, float c, float d, float e, float f, int symmetry) {
    import std.math : sin, cos;

    mut AffineTransform res;
    res.a = a; res.b = b; res.c = c; res.d = d; res.e = e; res.f = f;
    res.symmetry = symmetry;
    if (symmetry) {
        res.syms = new float mut[](symmetry);
        res.symc = new float mut[](symmetry);
        for (i in 0 .. symmetry) {
            float tau = 2 * 3.1415926538f;
            float angle = cast(int) i * tau / symmetry;
            res.syms[i] = sin(angle);
            res.symc[i] = cos(angle);
        }
    }
    return res;
}

float randvary(float v, XorShift rng, float f) {
    import std.math : sin;

    float interp;
    if (rng.rand % 4 != 0) {
        return v + sin(f * rng.randf) * rng.randf;
    }
    return v * (1 - f) + rng.randf * f;
}

struct AffineTransform
{
    float a, b, c, d, e, f;

    int symmetry;

    float mut[] syms, symc;

    int ii;

    Point apply(Point p) {
        mut auto p = Point(a * p.x + b * p.y + c, d * p.x + e * p.y + f);
        if (symmetry) {
            int i = ii % symmetry;
            ii = (ii + 1) % 65521;
            p = Point(p.x * symc[i] - p.y * syms[i], p.x * syms[i] + p.y * symc[i]);
        }
        return p;
    }

    void vary(XorShift rng, float f) {
        a = a.randvary(rng, f);
        b = b.randvary(rng, f);
        c = c.randvary(rng, f);
        d = d.randvary(rng, f);
        e = e.randvary(rng, f);
        this.f = this.f.randvary(rng, f);
    }
}

struct Function
{
    AffineTransform affine;

    int variation;

    Point apply(Point p) {
        import std.math : sin, cos;

        mut auto p = affine.apply(p);

        float r2() { return p.x * p.x + p.y * p.y; }
        float r() { return cast(float) sqrt(r2); }
        float angle() { return cast(float) atan2(p.y, p.x); }
        float pi = 3.1415926537f;
        int variation = variation;
        if (variation == 0) {
            return p;
        } else if (variation == 1) {
            return Point(sin(p.x), sin(p.y));
        } else if (variation == 2) {
            float f = 1/r2;
            return Point(f * p.x, f * p.y);
        } else if (variation == 3) {
            float sr = sin(r2), cr = cos(r2);
            return Point(p.x * sr - p.y * cr, p.x * cr + p.y * sr);
        } else if (variation == 4) {
            float dr = cast(float) (1 / sqrt(r2));
            return Point(dr * (p.x - p.y) * (p.x + p.y), dr * 2 * p.x * p.y);
        } else if (variation == 5) {
            return Point(angle / pi, r - 1);
        } else if (variation == 6) {
            float a = angle, r = r;
            return Point(r * sin(a + r), r * cos(a - r));
        } else if (variation == 7) {
            float a = angle, r = r;
            return Point(r * sin(a * r), -r * cos(a * r));
        } else if (variation == 8) {
            float f = angle / pi;
            return Point(f * sin(pi * r), f * cos(pi * r));
        } else if (variation == 9) {
            float a = angle, r = r, dr = 1 / r;
            return Point(dr * (cos(a) + sin(r)), dr * (sin(a) - cos(r)));
        } else if (variation == 10) {
            float a = angle, r = r;
            return Point(sin(a) / r, r * cos(a));
        } else if (variation == 11) {
            float a = angle, r = r;
            return Point(sin(a) * cos(r), cos(a) * sin(r));
        } else if (variation == 12) {
            float a = angle, r = r;
            float p0 = sin(a + r), p1 = cos(a - r);
            float p03 = p0 * p0 * p0, p13 = p1 * p1 * p1;
            return Point(r * (p03 + p13), r * (p03 - p13));
        } else if (variation == 13) {
            float sqr = cast(float) sqrt(r), a = angle;
            // TODO what is this
            float omik = 0;
            return Point(sqr * cos(a / 2 + omik), sqr * sin(angle / 2 + omik));
        } else if (variation == 14) {
            if (p.x > 0) {
                if (p.y > 0) return p;
                else return Point(p.x, p.y / 2);
            } else {
                if (p.y > 0) return Point(2 * p.x, p.y);
                else return Point(2 * p.x, p.y / 2);
            }
        }
    }

    void vary(XorShift rng, float f) {
        affine.vary(rng, f);
    }
}

struct MetaEntry
{
    float weight;

    Function f;

    void vary(XorShift rng, float f) {
        weight = weight.randvary(rng, f);
        this.f.vary(rng, f);
    }
}

struct Metafunction
{
    MetaEntry[] entries;

    AffineTransform affine;

    Color color;

    bool visible;

    Point apply(Point p) {
        Point result;
        for (entry in entries) {
            result.weightAdd(entry.f.apply(p), entry.weight);
        }
        return affine.apply(result);
    }

    void vary(XorShift rng, float f) {
        [entries[i].vary(rng, f) for i in 0 .. entries.length];
        affine.vary(rng, f);
    }
}

final class ChaosGame
{
    mut Metafunction[] functions;

    mut float[][] transitions; // [from][to], with each being the pass limit for randf to be under

    int rstate;

    this() { rstate = rand; }

    void initTransitions(XorShift rng) {
        float[] genTransitions() {
            float rescale(float low, float high, float f) {
                return low + (high - low) * f;
            }
            float eps = 0.005f;
            float[] base = [rescale(eps, 1 - eps, rng.randf()) for f in functions];
            float sum = [sum a for a in base];
            mut float[] result;
            mut float accum = 0;
            for (weight in base) {
                result ~= accum + weight / sum;
                accum += weight / sum;
            }
            return result;
        }
        transitions = [genTransitions for f in functions];
    }

    int selectTransition(XorShift rng, int from) {
        float f = rng.randf;
        for (mut int i = 0; i < functions.length; i += 1) {
            if (f < transitions[from][i]) return i;
        }
        return 0;
    }

    void apply(RangeCache cache, long l, long steps) {
        auto rng = new XorShift(l);
        mut Point p = Point(rng.randf, rng.randf);
        mut Color color = Color(0, 0, 0);
        mut int index = 0;
        for (long l in 0 .. steps) {
            index = selectTransition(rng, index);

            p = functions[index].apply(p);
            color = color.mix(functions[index].color, 0.5f);

            if (functions[index].visible) {
                cache.pset(p, color);
            }
        }
    }

    void vary(XorShift rng, float f) {
        // FIXME should be by reference
        [functions[i].vary(rng, f) for i in 0 .. functions.length];
        // if we allow the final to vary, it will visually dominate the output
        // finalFunction.vary(f);
    }
}

struct Color
{
    float r, g, b;

    Color mix(Color other, float factor)
    {
        return Color(
            r * (1 - factor) + other.r * factor,
            g * (1 - factor) + other.g * factor,
            b * (1 - factor) + other.b * factor);
    }
}

struct WeightedColor
{
    float r, g, b, i;
}

class Buffer
{
    int width; int height;

    mut long steps;

    WeightedColor mut[] back;

    this(this.width, this.height) {
        this.back = new WeightedColor mut[](width * height);
        for (i in 0 .. this.back.length)
            this.back[i] = WeightedColor(0, 0, 0, 0);
    }

    void writePixel(int offset, Color color) {
        auto ptr = &back[offset];
        ptr.r += color.r;
        ptr.g += color.g;
        ptr.b += color.b;
        ptr.i += 1;
    }

    void paint(SDL_Surface* surface) {
        // float adjf = 1 / powf(cast(int) (steps / (width * height)), 0.25);
        float adjf = 2.0f / cast(int) ((steps + width * height) / (width * height));
        for (y in 0 .. height) {
            for (x in 0 .. width) {
                auto color = this.back[y * width + x];
                float i = color.i * adjf;
                float i = logf(i + 1.01f) / (i + 0.01f);
                // float r = max(1 - color.r * i, 0);
                // float g = max(1 - color.g * i, 0);
                // float b = max(1 - color.b * i, 0);
                float r = min(color.r * i * adjf, 1.0f);
                float g = min(color.g * i * adjf, 1.0f);
                float b = min(color.b * i * adjf, 1.0f);
                int rgb =
                      (cast(int)(r * 0xff) & 0xff * 0x10000)
                    + (cast(int)(g * 0xff) & 0xff * 0x100)
                    + (cast(int)(b * 0xff) & 0xff * 0x1);
                // TODO direct write
                mut auto rect = SDL_Rect(cast(short) x, cast(short) y, 1, 1);
                SDL_FillRect(surface, &rect, rgb);
            }
        }
    }
}

float randnf(XorShift rng) {
    return rng.randf * 2.0f - 1.0f;
}

AffineTransform randAffine(XorShift rng) {
    mut int symmetry = 0;
    if (rng.irand % 10 == 0) {
        symmetry = rng.irand % 5 + 2;
    }
    vec2f vecInCircle() {
        while (true) {
            auto a = vec2f(rng.randnf, rng.randnf);
            if (a.length < 1) return a;
        }
    }
    // pi for luck
    float small(float x) => copysignf(powf(fabsf(x) * 1.1f, 3.141592f), x) / 1.1f;
    vec2f varyX = vecInCircle.(vec2f(x.small, y.small)) + vec2f(1, 0);
    vec2f varyY = vecInCircle.(vec2f(x.small, y.small)) + vec2f(0, 1);
    return makeAffineTransform(
        a=varyX.x, b=varyY.x, c=rng.randnf.small,
        d=varyX.y, e=varyY.y, f=rng.randnf.small,
        symmetry);
}

Function randFun(XorShift rng) {
    return Function(randAffine(rng), rng.irand % 15);
}

Metafunction randMetafun(XorShift rng, bool visible) {
    mut MetaEntry mut[] entries;
    mut float weightsum;
    for (i in 0 .. rng.irand % 4 + 2) {
        float weight = rng.randf;
        entries ~= MetaEntry(weight, rng.randFun);
        weightsum += weight;
    }
    for (i in 0 .. entries.length) {
        entries[i].weight /= weightsum;
    }
    return Metafunction(entries.freeze, rng.randAffine, Color(rng.randf, rng.randf, rng.randf), visible);
}

class RenderGame : Task
{
    ChaosGame game;
    RangeCache cache;
    Buffer buffer;
    long index;
    long steps;

    this(this.game, this.buffer, this.index, this.steps) {
        this.cache = new RangeCache(buffer);
    }

    override void run() {
        game.apply(cache, index, steps);
        cache.flush;
    }
}

struct CacheEntry
{
    Color color;
    int offset;
}

class RangeCache
{
    Buffer buffer;

    bool cache;

    int width, height;
    int sections;
    int sectionLength;
    int sectionCovers;
    CacheEntry mut[] ranges;
    int mut[] offsets;
    float ratio;
    float hwidth, hheight;

    mut long steps;

    this(this.buffer) {
        this.width = buffer.width;
        this.height = buffer.height;
        this.ratio = height * 1.0f / width;
        this.hwidth = width / 2;
        this.hheight = height / 2;
        // TODO how to balance
        if (true || width * height < 4194304) {
            return;
        }
        this.cache = true;
        this.sectionCovers = 16384;
        print("w " ~ itoa(width) ~ " h " ~ itoa(height) ~ " covers 16384?");
        assert(width * height % sectionCovers == 0);
        this.sections = width * height / sectionCovers;
        this.offsets = new int mut[](sections);
        this.ranges = new CacheEntry mut[](this.sections * 32);
        // this.sectionLength = ranges.length / sections;
        // print("range " ~ itoa(cast(int) this.ranges.length) ~ " over " ~ itoa(this.sections));
        assert(this.ranges.length % this.sections == 0);
        this.sectionLength = cast(int) (this.ranges.length / this.sections);
    }

    void pset(Point p, Color color) {
        this.steps += 1;

        int x = cast(int)((p.x * ratio + 1) * hwidth), y = cast(int)((p.y + 1) * hheight);
        if (x < 0 || x >= width || y < 0 || y >= height) return;
        int offset = y * width + x;

        // Does this actually do anything ... ?
        if (!this.cache) {
            this.buffer.writePixel(offset, color);
            return;
        }

        int section = offset / 16384; // TODO right shift
        if (offsets[section] == sectionLength)
        {
            flushSection(section);
        }
        ranges[section * sectionLength + offsets[section]] = CacheEntry(color, offset);
        offsets[section] += 1;
    }
    void flushSection(int section) {
        for (i in 0 .. offsets[section]) {
            auto entry = ranges[section * sectionLength + i];
            buffer.writePixel(entry.offset, entry.color);
        }
        offsets[section] = 0;
    }
    void flush() {
        buffer.steps += this.steps;
        for (i in 0 .. this.sections) {
            flushSection(cast(int) i);
        }
    }
}

extern(C) void setbuf(void* stream, char* buf);

void main(mut string[] args) {
    import std.string : atoi;

    string executable = args[0];
    args = args[1 .. $];
    mut int seed = 1;
    mut long iters = 1000;
    mut int width = 640, height = 480;
    mut string output = "out.png";
    mut bool htmlProgress;
    mut float varyFactor = 0;
    for (mut size_t i = 0; i < args.length; i += 1) {
        auto arg = args[i];
        if (arg == "-h" || arg == "--help") {
            print("Usage: ./fflame
    -e, --seed num        : main rng seed. Specifies the image.
    -i, --iters n         : number of iterations before exiting
    -s, --size wxh        : target image size
    -o, --output file.png : png file to save to
    -x, --html            : print progress as html/js progress bar
    -v, --vary f          : randomize and interpolate (for animations)
    -h, --help            : this screen");
            return;
        }
        if (arg == "-e" || arg == "--seed") {
            seed = args[i + 1].atoi;
            i += 1;
            continue;
        }
        if (arg == "-i" || arg == "--iters") {
            iters = args[i + 1].toStringz.atol;
            i += 1;
            continue;
        }
        if (arg == "-s" || arg == "--size") {
            string size = args[i + 1];
            auto wh = size.split("x");
            width = wh[0].atoi;
            height = wh[1].atoi;
            i += 1;
            continue;
        }
        if (arg == "-x" || arg == "--html") {
            htmlProgress = true;
            continue;
        }
        if (arg == "-v" || arg == "--vary") {
            import std.string : atof;

            varyFactor = cast(float) args[i + 1].atof;
            i += 1;
            continue;
        }
        if (arg == "-o" || arg == "--output") {
            output = args[i + 1];
            i += 1;
            continue;
        }
        print("Unknown argument " ~ arg);
        assert(false);
    }
    auto rng = new XorShift(seed);
    mut int subtasks = 100;
    if (width < 400 && height < 400) subtasks = 4;

    auto pool = new ThreadPool(7);

    auto game = new ChaosGame;
    int limit = 1 + rng.irand % 10;
    int visible = rng.irand % limit;
    for (i in 0 .. limit)
        game.functions ~= rng.randMetafun(i == visible);
    game.initTransitions(rng);
    game.vary(rng, varyFactor);

    Buffer buffer = new Buffer(width, height);

    for (j in 0 .. subtasks) {
        auto task = new RenderGame(game, buffer, j, iters / subtasks);
        pool.addTask(task);
    }

    void printProgressBar(float progress) {
        if (!htmlProgress) return;
        int percent = cast(int) (progress * 100);
        print("<script type=\"text/javascript\">\$(\"#progress\").html(\""
            ~ itoa(percent)
            ~ "%\");</script>");
    }
    setbuf(neat_runtime_stdout(), null);
    pool.waitComplete(&printProgressBar);
    auto surface = SDL_CreateRGBSurface(0, width, height, 32, 0, 0, 0, 0);
    buffer.paint(surface);
    mut SDL_RWops* out_;
    if (output == "-") {
        out_ = SDL_RWFromFP(cast(_IO_FILE*) neat_runtime_stdout(), false);
    } else {
        out_ = SDL_RWFromFile(toStringz(output), toStringz("wb"));
    }
    IMG_SavePNG_RW(surface, out_, true);
}