#!/usr/bin/env python3
"""Compare world seeds by their island's biomes, without generating chunks.

Boots the test server once (world "seed-search"), runs
`/cobblesurvival biomepreview <seed> <radius> <step>` for every candidate, and
scores each island:

  * coverage: every surface biome from world/biomes.json reaches its target
    share of the land (TARGETS below);
  * patches: few, large patches per biome instead of many crumbs;
  * balance: no biome takes more than MAX_SHARE of the land;
  * start: the land around 0,0 is gentle (plains, forest, cherry grove);
  * caves: lush caves, dripstone caves, deep dark and basalt deltas all occur
    under land.

Writes build/seed-search/report.json and candidates.png (best seeds side by
side). The coastline in the preview is approximate (see BiomePreview.kt).

Usage: tools/seed_search.py [--count N] [--prefix P] [--apply]
       --apply writes the best seed into world/world.json
"""
import argparse
import json
import shutil
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from mcserver import BUILD, ROOT, SERVER, Server, eula_accepted, prepare  # noqa: E402
from world_report import colour  # noqa: E402

LEVEL = "seed-search"
OUT = BUILD / "seed-search"
RADIUS, STEP = 1700, 24
SURFACE = ["minecraft:plains", "minecraft:forest", "minecraft:cherry_grove", "minecraft:mangrove_swamp",
           "minecraft:jungle", "minecraft:desert", "minecraft:badlands", "minecraft:windswept_hills",
           "minecraft:jagged_peaks"]
TARGETS = {b: 0.05 for b in SURFACE} | {"minecraft:cherry_grove": 0.03, "minecraft:mangrove_swamp": 0.03,
                                        "minecraft:jagged_peaks": 0.03, "minecraft:windswept_hills": 0.03}
CAVES = ["minecraft:lush_caves", "minecraft:dripstone_caves", "minecraft:deep_dark", "minecraft:basalt_deltas"]
GENTLE = {"minecraft:plains", "minecraft:forest", "minecraft:cherry_grove"}
MAX_SHARE = 0.35
START_RADIUS = 350


def patches(mask):
    """Number of 4-connected patches of at least 3 cells."""
    seen = np.zeros_like(mask)
    count = 0
    for start in zip(*np.nonzero(mask)):
        if seen[start]:
            continue
        seen[start] = True
        stack, size = [start], 0
        while stack:
            z, x = stack.pop()
            size += 1
            for nz, nx in ((z + 1, x), (z - 1, x), (z, x + 1), (z, x - 1)):
                if 0 <= nz < mask.shape[0] and 0 <= nx < mask.shape[1] and mask[nz, nx] and not seen[nz, nx]:
                    seen[nz, nx] = True
                    stack.append((nz, nx))
        count += size >= 3
    return count


def score(preview):
    n = preview["size"]
    names = preview["biomes"]
    surface = np.array(preview["surface"]).reshape(n, n)
    land = np.array(preview["height"]).reshape(n, n) >= preview["sea_level"]
    total = land.sum() or 1
    shares, pieces = {}, {}
    for b in SURFACE:
        mask = land & (surface == names.index(b)) if b in names else np.zeros_like(land)
        shares[b] = float(mask.sum() / total)
        pieces[b] = patches(mask)
    coverage = sum(min(shares[b] / TARGETS[b], 1.0) for b in SURFACE)
    fragmentation = sum(max(p - 1, 0) for p in pieces.values()) / len(SURFACE)
    balance = max(0.0, max(shares.values()) - MAX_SHARE) * 10
    coords = (np.arange(n) * preview["step"] - preview["radius"])
    near = (np.hypot(*np.meshgrid(coords, coords)) < START_RADIUS) & land
    gentle_ids = [names.index(b) for b in GENTLE if b in names]
    start = float(np.isin(surface[near], gentle_ids).mean()) if near.any() else 0.0
    caves = {b: 0 for b in CAVES}
    for grid in preview["caves"].values():
        layer = np.array(grid).reshape(n, n)[land]
        for b in CAVES:
            if b in names:
                caves[b] += int((layer == names.index(b)).sum())
    complete = all(shares[b] >= 0.01 for b in SURFACE) and all(caves.values())
    total_score = coverage + 2 * start - 0.15 * fragmentation - balance + (5 if complete else 0)
    return {"seed": preview["seed"], "score": round(total_score, 3), "complete": complete,
            "land_cells": int(total), "start_gentle": round(start, 2),
            "shares": {b.split(":")[1]: round(v, 3) for b, v in shares.items()},
            "patches": {b.split(":")[1]: v for b, v in pieces.items()},
            "cave_cells": {b.split(":")[1]: v for b, v in caves.items()}}


def tile(preview, label, size=300):
    from PIL import Image, ImageDraw
    n = preview["size"]
    surface = np.array(preview["surface"]).reshape(n, n)
    height = np.array(preview["height"]).reshape(n, n)
    lut = np.array([colour(b) for b in preview["biomes"]], dtype=np.uint8)
    rgb = lut[surface]
    rgb[height < preview["sea_level"]] = (20, 50, 120)
    img = Image.fromarray(rgb).resize((size, size), Image.NEAREST)
    ImageDraw.Draw(img).text((6, 6), label, fill=(255, 255, 255), stroke_width=2, stroke_fill=(0, 0, 0))
    return img


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--count", type=int, default=60)
    parser.add_argument("--prefix", default="isle")
    parser.add_argument("--apply", action="store_true")
    args = parser.parse_args()
    if not eula_accepted():
        sys.exit("build/server/eula.txt does not accept the Minecraft EULA yet; run tools/smoke_test.py --accept-eula first.")

    recipe = json.loads((ROOT / "world/world.json").read_text())
    seeds = [recipe["seed"]] + [f"{args.prefix}-{i}" for i in range(1, args.count + 1)]
    args_file = prepare()
    shutil.rmtree(SERVER / LEVEL, ignore_errors=True)
    shutil.rmtree(OUT, ignore_errors=True)
    OUT.mkdir(parents=True)
    results, previews = [], {}
    props = {"level-name": LEVEL, "level-seed": recipe["seed"], "level-type": recipe["level_type"]}
    with Server(args_file, props, stdout_name="seed-search-stdout.log") as server:
        server.wait_ready()
        for i, seed in enumerate(seeds, 1):
            started = time.time()
            reply = server.command(f"cobblesurvival biomepreview {seed} {RADIUS} {STEP}")
            if "written" not in reply:
                raise RuntimeError(f"preview failed for {seed}: {reply}")
            preview = json.loads((SERVER / LEVEL / "cobblesurvival/biome_preview" / f"{seed}.json").read_text())
            previews[seed] = preview
            results.append(score(preview))
            r = results[-1]
            print(f"[{i}/{len(seeds)}] {seed:16} score {r['score']:6.2f} complete={r['complete']!s:5} "
                  f"start={r['start_gentle']:.2f} ({time.time() - started:.1f}s)")
        server.stop()

    results.sort(key=lambda r: -r["score"])
    (OUT / "report.json").write_text(json.dumps(results, indent=2) + "\n")
    from PIL import Image
    best = results[:8]
    sheet = Image.new("RGB", (4 * 300, 2 * 300), (24, 24, 24))
    for k, r in enumerate(best):
        sheet.paste(tile(previews[r["seed"]], f"#{k + 1} {r['seed']}  {r['score']}"), ((k % 4) * 300, (k // 4) * 300))
    sheet.save(OUT / "candidates.png")
    print(f"\nbest: {best[0]['seed']} (score {best[0]['score']}); report: {(OUT / 'report.json').relative_to(ROOT)}")
    print(json.dumps(best[0], indent=2))
    if args.apply:
        recipe["seed"] = best[0]["seed"]
        (ROOT / "world/world.json").write_text(json.dumps(recipe, indent=2) + "\n")
        print(f"world/world.json seed set to {best[0]['seed']}")


if __name__ == "__main__":
    main()