# Copyright (C) 2022 RozK
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from os import urandom
from operator import mul
from itertools import product, combinations, chain, starmap, repeat
from math import prod, sqrt, floor, ceil, dist, hypot, pi, cos, sin
from random import seed as set_seed, random, uniform, triangular, sample
from array import array
from struct import unpack

from game.math import vec3_normal_rgb10a2

class Generator:
    def __init__(self, size):
        seed = unpack('I', urandom(4))[0] # 666 # 2749145609 #
        print("Seed =", seed)
        set_seed(seed)
        self.size = size
        self.map_coords = tuple(product(range(size), repeat = 2))
        self.field_coords = tuple(starmap(lambda y, x: (y + 0.5, x + 0.5), self.map_coords))
        self.generate_field()
        self.generate_heights()
        self.generate_rivers()
        self.smooth_heights()
        self.generate_normals()
        self.smooth_normals()
        self.pack()

    def weighted_samples(self, radius):
        samples = []
        smin = floor(radius)
        tw = 0.0
        for y, x in product(range(-smin, smin + 1), repeat = 2):
            w = 1.0 - hypot(y, x) / radius
            if w > 0.0:
                tw += w
                samples.append([y, x, w])
        iw = 1.0 / tw
        for sample in samples:
            sample[2] *= iw
        return samples

    def generate_field(self):
        half = self.size // 2
        vhalf = (half, half)

        rmin = self.size / 16
        rmax = self.size / 8
        vrad = self.size / 32
        cones = [(vhalf, vrad, 0.2)]
        for _ in range(180):
            r = uniform(rmin, rmax)
            d = uniform(r, half + r)
            a = uniform(-pi, pi)
            p = d / half
            cones.append(((d * cos(a) + half, d * sin(a) + half), r, p))

        def height(_position):
            def influence(_center, _radius, _power):
                d = dist(_position, _center)
                return ((_radius ** 2.0) / (d ** 2.0) if d > _radius else (d ** 2.0) / (_radius ** 2.0)) * _power
            return sum(starmap(influence, cones))

        heights = list(map(height, self.field_coords))
        heights = list(map(mul, heights, repeat(1.0 / max(heights))))

        self.volcano_c = vhalf
        self.volcano_r = vrad
        self.cones = cones
        self.heights = heights

    def generate_heights(self):
        size = self.size
        half = size // 2
        vhalf = (half, half)
        heights = self.heights

        def shape(_position, _height):
            d = dist(vhalf, _position) / half
            return max(0.0, _height - d) ** 2.5
        heights = list(map(shape, self.field_coords, heights))

        vy, vx = self.volcano_c
        vrad = self.volcano_r
        vmin = ceil(vrad)
        vmax = vrad - 2.0
        addrs = []
        hmin = 1.0
        for y, x in product(range(-vmin, vmin), repeat = 2):
            if hypot(y + 0.5, x + 0.5) < vmax:
                addr = (y + vy) * size + (x + vx)
                addrs.append(addr)
                hmin = min(hmin, heights[addr])
        hmin -= 8.0 / 255.0
        for addr in addrs:
            heights[addr] = hmin

        self.heights = list(map(mul, heights, repeat(255.0 / max(heights))))

    def generate_rivers(self):
        size = self.size
        heights = self.heights

        cw = 1.0 / sqrt(2.0)
        offsets = (
            # y   x   w      y  x   w      y  x   w
            (-1, -1,  cw), (-1, 0, 1.0), (-1, 1,  cw),
            ( 0, -1, 1.0),               ( 0, 1, 1.0),
            ( 1, -1,  cw), ( 1, 0, 1.0), ( 1, 1,  cw))

        def river(_ry, _rx):
            path = []
            mins = []
            minh = heights[_ry * size + _rx]
            flowing = True
            while flowing:
                flowing = False
                path.append((_ry, _rx))
                mins.append(minh)
                rh = heights[_ry * size + _rx]
                if rh == 0.0:
                    break
                for by, bx in reversed(path):
                    ns = []
                    for oy, ox, w in offsets:
                        ny = (_ry + oy) % size
                        nx = (_rx + ox) % size
                        if (ny, nx) not in path:
                            nh = heights[ny * size + nx]
                            ns.append(((nh - rh) * w, min(minh, nh), ny, nx))
                    if ns:
                        _, minh, _ry, _rx = min(ns)
                        flowing = True
                        break
            return [(rx, ry, rh) for (rx, ry), rh in zip(path, mins)]

        def center(_cy, _cx):
            return (heights[_cy * size + _cx], _cy, _cx)
        centers = sorted(
            [center(round(cy) % size, round(cx) % size) for (cy, cx), _, _ in self.cones[1:]],
            reverse = True)
        sources = [(sy, sx) for _, sy, sx in centers[:32]]

        paths = []
        for sy, sx in sources:
            paths.append(river(sy, sx))

        rivers = [0.0 for _ in range(size ** 2)]
        path = max(paths, key = len)
        for ry, rx, minh in path:
            minh = max(0.0, minh - 1.0)
            addr = ry * size + rx
            heights[addr] = minh
            for oy, ox, _ in offsets:
                heights[(ry + oy) * size + (rx + ox)] = minh
            rivers[addr] = 1.0
        self.rivers = rivers

    def smooth_heights(self):
        size = self.size
        heights = self.heights

        samples = self.weighted_samples(2.0)
        def smooth(_y, _x):
            return sum(heights[((_y + dy) % size) * size + ((_x + dx) % size)] * w for dy, dx, w in samples)
        heights = list(starmap(smooth, self.map_coords))

        self.heights = list(map(mul, heights, repeat(255.0 / max(heights))))

    def generate_normals(self):
        size = self.size
        heights = self.heights

        def normal(_y, _x):
            height = heights[_y * size + _x]
            def direction(_dy, _dx):
                dz = (heights[((_y - _dy) % size) * size + ((_x + _dx) % size)] - height) / 8.0
                di = 1.0 / sqrt(_dx ** 2.0 + _dy ** 2.0 + dz ** 2.0) #OPTIM: dx² + dy² = 1.0
                return (_dx * di, _dy * di, dz * di)
            hx, hy, hz = direction(0, 1)
            vx, vy, vz = direction(1, 0)
            return (hy * vz - hz * vy, hz * vx - hx * vz, hx * vy - hy * vx)

        self.normals = list(starmap(normal, self.map_coords))

    def smooth_normals(self):
        size = self.size
        normals = self.normals

        samples = self.weighted_samples(2.0)
        def smooth(_y, _x):
            tx = ty = tz = 0.0
            for dy, dx, w in samples:
                nx, ny, nz = normals[((_y + dy) % size) * size + ((_x + dx) % size)]
                tx += nx * w
                ty += ny * w
                tz += nz * w
            di = 1.0 / sqrt(tx ** 2.0 + ty ** 2.0 + tz ** 2.0)
            return (tx * di, ty * di, tz * di)

        self.normals = list(starmap(smooth, self.map_coords))

    def pack(self):
        self.packed_heights = array('f', self.heights)
        self.packed_normals = array('I', tuple(map(vec3_normal_rgb10a2, self.normals)))

    def unpack(self, y, x):
        addr = y * self.size + x
        nx, ny, nz = self.normals[addr]
        return (nx, ny, nz, self.heights[addr])
