#!/usr/bin/env python3
"""Generates the synthetic input rasters of the Landscape Program (ЛП) test.

Writes 13 grid-aligned GeoTIFFs, one per input slot (file name = slot code),
plus one deliberately misaligned raster for the negative test (L13):

    <out>/lp-z-biodiversity.tif ... <out>/lp-p-tourism.tif   (13 files)
    <out>/misaligned-lp-u-veg.tif                            (half size, 2x pixel)

Every input raster: int16, EPSG:32642 (UTM 42N), 100 m pixels, default
3000 x 2000 pixels, classes 1..5 from smooth trigonometric patterns (a
different pattern per slot, so ties and every class combination occur), and a
NoData (-9999) frame 40 pixels wide. Extra edge cases exercise the pixel rules:

  * lp-u-crop: a 200 x 200 NoData block — U is computed from water, soil and
    veg only (the crop/veg union rule of u_integrated / u_max);
  * lp-u-veg: a 100 x 100 NoData block overlapping half of the crop block —
    where both crop and veg are missing, U must be NoData;
  * lp-p-snow: a 120 x 120 block of 0 and a 60 x 60 block of 6 — values that
    are not NoData but outside 1..5, hence invalid (the "valid iff 1 <= v <= 5"
    rule).

Output is deterministic (no random numbers), so the same command always
produces byte-identical pixel values.

    python lp_gen.py --out /tmp/lp-inputs [--width 3000 --height 2000]
"""
import argparse
import os
import sys

import numpy as np
import rasterio
from rasterio.transform import from_origin

SLOTS = [
    "lp-z-biodiversity", "lp-z-lcr", "lp-z-soil-protection", "lp-z-wr",
    "lp-u-surf-water", "lp-u-soil", "lp-u-crop", "lp-u-veg",
    "lp-p-irrigation", "lp-p-snow", "lp-p-crop", "lp-p-pasture", "lp-p-tourism",
]
NODATA = -9999
CRS = "EPSG:32642"
ORIGIN_X, ORIGIN_Y = 300000.0, 6100000.0  # UTM 42N, North Kazakhstan region
PIXEL = 100.0
FRAME = 40


def pattern(k, width, height):
    """Smooth class field 1..5 for slot index k."""
    y, x = np.mgrid[0:height, 0:width].astype(np.float64)
    fx = 2 * np.pi * (1.0 + 0.37 * k) / width
    fy = 2 * np.pi * (0.8 + 0.29 * ((k * 5) % 7)) / height
    v = (np.sin(fx * x + 0.9 * k) * np.cos(fy * y + 0.4 * k)
         + 0.5 * np.sin(fx * 2.3 * x + fy * 1.7 * y + k))
    v = (v + 1.5) / 3.0  # ~0..1
    cls = np.floor(np.clip(v, 0, 0.999999) * 5).astype(np.int16) + 1
    return cls


def block(arr, r0, c0, h, w, value):
    arr[r0:r0 + h, c0:c0 + w] = value


def write(path, arr, transform):
    profile = dict(driver="GTiff", width=arr.shape[1], height=arr.shape[0], count=1,
                   dtype="int16", crs=CRS, transform=transform, nodata=NODATA,
                   tiled=True, blockxsize=256, blockysize=256, compress="DEFLATE")
    with rasterio.open(path, "w", **profile) as dst:
        dst.write(arr, 1)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--width", type=int, default=3000)
    ap.add_argument("--height", type=int, default=2000)
    args = ap.parse_args()
    w, h = args.width, args.height
    if w < 400 or h < 400:
        sys.exit("width/height must be >= 400")
    os.makedirs(args.out, exist_ok=True)
    transform = from_origin(ORIGIN_X, ORIGIN_Y, PIXEL, PIXEL)

    for k, slot in enumerate(SLOTS):
        arr = pattern(k, w, h)
        arr[:FRAME, :] = NODATA
        arr[-FRAME:, :] = NODATA
        arr[:, :FRAME] = NODATA
        arr[:, -FRAME:] = NODATA
        cr, cc = h // 2 - 100, w // 2 - 100  # centre block anchor
        if slot == "lp-u-crop":
            block(arr, cr, cc, 200, 200, NODATA)
        if slot == "lp-u-veg":
            block(arr, cr + 50, cc + 150, 100, 100, NODATA)
        if slot == "lp-p-snow":
            block(arr, 300, 300, 120, 120, 0)
            block(arr, 300, 500, 60, 60, 6)
        write(os.path.join(args.out, slot + ".tif"), arr, transform)
        print("wrote", slot + ".tif")

    # Misaligned raster for L13: half the size, double pixel => same footprint,
    # different grid.
    mis = pattern(7, w // 2, h // 2)
    write(os.path.join(args.out, "misaligned-lp-u-veg.tif"), mis,
          from_origin(ORIGIN_X, ORIGIN_Y, PIXEL * 2, PIXEL * 2))
    print("wrote misaligned-lp-u-veg.tif (%dx%d, %d m)" % (w // 2, h // 2, int(PIXEL * 2)))
    return 0


if __name__ == "__main__":
    sys.exit(main())