"""Reading, writing and resampling heightmaps with nothing but numpy, so the authoring scripts run on the
engine's own Python (which has no PIL). Greyscale PNG in 8 or 16 bit, raw 16-bit little-endian (.r16 / .raw,
what World Machine, Gaea and the engine's own exporter write), bilinear resampling and a centred square crop:
enough to take a real heightmap from any of the usual sources and put it on the landscape.
"""
import math
import struct
import zlib

import numpy as np

PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


def write_png(path, data):
    """Greyscale PNG, 8 or 16 bit from the array dtype. Row filter 0, one zlib stream."""
    if data.dtype == np.uint16:
        depth, payload = 16, data.astype(">u2")
    else:
        depth, payload = 8, data.astype(np.uint8)
    height, width = data.shape
    raw = b"".join(b"\x00" + payload[y].tobytes() for y in range(height))

    def chunk(kind, body):
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body) & 0xFFFFFFFF)

    ihdr = struct.pack(">IIBBBBB", width, height, depth, 0, 0, 0, 0)
    with open(path, "wb") as f:
        f.write(PNG_SIGNATURE + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw, 6)) + chunk(b"IEND", b""))


def read_png_header(path):
    """(width, height, bit_depth, colour_type) without decoding the image."""
    with open(path, "rb") as f:
        head = f.read(8 + 8 + 13)
    if head[:8] != PNG_SIGNATURE or head[12:16] != b"IHDR":
        raise ValueError(f"{path}: not a PNG")
    width, height, depth, colour_type = struct.unpack(">IIBB", head[16:26])
    return width, height, depth, colour_type


def _unfilter_sequential(filter_type, row, prev, bpp):
    """Average and Paeth depend on the byte just decoded, so they go pixel by pixel. Rare in practice; a
    4081x4081 16-bit file with every row Paeth-filtered takes some tens of seconds, once, on import."""
    out = bytearray(row)
    n = len(out)
    if filter_type == 3:
        for i in range(n):
            left = out[i - bpp] if i >= bpp else 0
            out[i] = (out[i] + ((left + prev[i]) >> 1)) & 0xFF
    else:
        for i in range(n):
            if i >= bpp:
                a, c = out[i - bpp], prev[i - bpp]
            else:
                a, c = 0, 0
            b = prev[i]
            p = a + b - c
            pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
            if pa <= pb and pa <= pc:
                predictor = a
            elif pb <= pc:
                predictor = b
            else:
                predictor = c
            out[i] = (out[i] + predictor) & 0xFF
    return out


def read_png(path):
    """The first channel of a non-interlaced PNG as a 2D uint8 or uint16 array (greyscale, grey+alpha, RGB
    and RGBA are accepted; palette and interlaced files are not)."""
    with open(path, "rb") as f:
        blob = f.read()
    if blob[:8] != PNG_SIGNATURE:
        raise ValueError(f"{path}: not a PNG")
    pos, idat, ihdr = 8, [], None
    while pos + 8 <= len(blob):
        length, kind = struct.unpack(">I4s", blob[pos:pos + 8])
        body = blob[pos + 8:pos + 8 + length]
        pos += 12 + length
        if kind == b"IHDR":
            ihdr = struct.unpack(">IIBBBBB", body)
        elif kind == b"IDAT":
            idat.append(body)
        elif kind == b"IEND":
            break
    if ihdr is None:
        raise ValueError(f"{path}: no IHDR")
    width, height, depth, colour_type, _, _, interlace = ihdr
    channels = {0: 1, 2: 3, 4: 2, 6: 4}.get(colour_type)
    if channels is None or depth not in (8, 16) or interlace != 0:
        raise ValueError(f"{path}: unsupported PNG (colour type {colour_type}, {depth} bit, interlace {interlace}); "
                         "use a non-interlaced 8 or 16 bit greyscale or RGB file")
    bytes_per_sample = depth // 8
    bpp = channels * bytes_per_sample
    stride = width * bpp
    data = zlib.decompress(b"".join(idat))
    if len(data) != height * (stride + 1):
        raise ValueError(f"{path}: PNG data is {len(data)} bytes, expected {height * (stride + 1)}")

    rows = np.empty((height, stride), dtype=np.uint8)
    prev = np.zeros(stride, dtype=np.uint8)
    for y in range(height):
        start = y * (stride + 1)
        filter_type = data[start]
        row = np.frombuffer(data, dtype=np.uint8, count=stride, offset=start + 1)
        if filter_type == 0:
            out = row.copy()
        elif filter_type == 1:
            out = (np.cumsum(row.reshape(width, bpp), axis=0, dtype=np.uint64) & 0xFF).astype(np.uint8).reshape(stride)
        elif filter_type == 2:
            out = ((row.astype(np.uint16) + prev) & 0xFF).astype(np.uint8)
        elif filter_type in (3, 4):
            out = np.frombuffer(bytes(_unfilter_sequential(filter_type, bytes(row), bytes(prev), bpp)), dtype=np.uint8)
        else:
            raise ValueError(f"{path}: bad PNG filter {filter_type} on row {y}")
        rows[y] = out
        prev = rows[y]

    dtype = ">u2" if depth == 16 else np.uint8
    samples = rows.reshape(height, width * channels * bytes_per_sample).view(dtype).reshape(height, width, channels)
    first = samples[:, :, 0]
    return first.astype(np.uint16) if depth == 16 else first.astype(np.uint8)


def read_r16(path, width=None):
    """Raw 16-bit little-endian samples, square unless a width is given."""
    values = np.fromfile(path, dtype="<u2")
    if width is None:
        width = math.isqrt(len(values))
        if width * width != len(values):
            raise ValueError(f"{path}: {len(values)} samples is not a square; give the width in the manifest")
    if len(values) % width != 0:
        raise ValueError(f"{path}: {len(values)} samples do not divide by width {width}")
    return values.reshape(len(values) // width, width).astype(np.uint16)


def read_heightmap(path, width=None):
    """Any supported file as a 2D uint16 array with the full 0..65535 range (8-bit files are widened)."""
    lower = path.lower()
    if lower.endswith(".png"):
        values = read_png(path)
        return values.astype(np.uint16) * 257 if values.dtype == np.uint8 else values
    if lower.endswith((".r16", ".raw")):
        return read_r16(path, width)
    raise ValueError(f"{path}: unknown heightmap format; use 16-bit PNG or raw .r16")


def center_crop_square(values):
    height, width = values.shape
    side = min(height, width)
    y0, x0 = (height - side) // 2, (width - side) // 2
    return values[y0:y0 + side, x0:x0 + side]


def block_mean(values, factor):
    """Downsample by an integer factor with a box filter, trimming the edge that does not divide."""
    height, width = values.shape
    height, width = height // factor * factor, width // factor * factor
    trimmed = values[:height, :width].astype(np.float32)
    return trimmed.reshape(height // factor, factor, width // factor, factor).mean(axis=(1, 3))


def resample(values, size):
    """Bilinear resample of a 2D array to size x size, box-filtered first when shrinking by 2x or more."""
    source = values.astype(np.float32)
    factor = min(source.shape) // size
    if factor >= 2:
        source = block_mean(source, factor)
    src_h, src_w = source.shape
    if (src_h, src_w) == (size, size):
        return source
    ys = np.linspace(0.0, src_h - 1, size, dtype=np.float32)
    xs = np.linspace(0.0, src_w - 1, size, dtype=np.float32)
    y0 = np.floor(ys).astype(np.int64)
    x0 = np.floor(xs).astype(np.int64)
    y1 = np.minimum(y0 + 1, src_h - 1)
    x1 = np.minimum(x0 + 1, src_w - 1)
    ty = (ys - y0)[:, None]
    tx = (xs - x0)[None, :]
    top = source[np.ix_(y0, x0)] * (1 - tx) + source[np.ix_(y0, x1)] * tx
    bottom = source[np.ix_(y1, x0)] * (1 - tx) + source[np.ix_(y1, x1)] * tx
    return (top * (1 - ty) + bottom * ty).astype(np.float32)