getJsonBases.py

# getJsonBases.py
# Génère un fichier JSON par tuile 5°x5° (ex: E45N50.json) contenant toutes les bases
# dont le centre tombe dans la zone correspondante. Gestion robuste de l'écriture
# pour éviter PermissionError sous Windows (retries + fallback).

import os
import re
import csv
import json
import math
import hashlib
import tempfile
import time
import shutil
from pathlib import Path
from collections import Counter
import numpy as np
import rasterio
from rasterio.windows import Window
from tqdm import tqdm

# ---------------- CONFIG ----------------
TILES_DIR = "data"           # dossier contenant les .tif
OUTPUT_SUMMARY = "all_bases_index_summary.json"
CHECKPOINT_CSV = "bases_progress.csv"  # fichier de contrôle / reprise
PIX_SIDE = 1700
BAND_TO_PREFER = None   # None -> prefer last band (ds.count), else integer
OCEAN_CODE = 210
TOPK = 3
NODATA_DEFAULTS = {0}
JSON_OUT_DIR = "json"   # dossier de sortie pour les JSON par tuile (ex: json/E45N50.json)
REPLACE_RETRIES = 3     # nombre de tentatives pour os.replace avant fallback
REPLACE_SLEEP = 0.2     # secondes entre tentatives
# ----------------------------------------

places = [
    'Fort', 'Camp', 'Base', 'Bastion', 'Citadel',
    'Barracks', 'Outpost', 'Bunker', 'Shelter', 'Hangar',
    'Station', 'Center', 'Fortlet', 'Depot', 'Encampment',
    'Battery', 'Installation', 'Redoubt', 'Fortress', 'Casemate',
    'Post', 'Airbase', 'Prison',
]
animals = [
    'Scorpion', 'Platypus', 'Falcon', 'Lion', 'Panda',
    'Eagle', 'Narwhal', 'Tiger', 'Phoenix', 'Spider',
    'Turtle', 'Dog', 'Raven', 'Wolf', 'Jaguar',
    'Cobra', 'Hawk', 'Leopard', 'Anaconda', 'Panther',
    'Lynx', 'Bear', 'Bull', 'Shark', 'Monkey',
    'Elephant', 'Hornet', 'Crab',
]
colors = [
    'Green', 'Blue', 'Black', 'Purple', 'Gold',
    'Silver', 'Copper', 'Brown', 'Orange', 'Red',
    'Yellow', 'Gray', 'White',
]

LC = {
    10: ("Rainfed cropland", (255,255,100)), 11: ("Herbaceous cropland", (255,255,100)),
    12: ("Tree/shrub cropland", (255,255,0)), 20: ("Irrigated cropland", (170,240,240)),
    51: ("Open evergreen broadleaved forest", (76,115,0)), 52: ("Closed evergreen broadleaved forest", (0,100,0)),
    61: ("Open deciduous broadleaved forest", (170,200,0)), 62: ("Closed deciduous broadleaved forest", (0,160,0)),
    71: ("Open evergreen needle-leaved forest", (0,80,0)), 72: ("Closed evergreen needle-leaved forest", (0,60,0)),
    81: ("Open deciduous needle-leaved forest", (40,100,0)), 82: ("Closed deciduous needle-leaved forest", (40,80,0)),
    91: ("Open mixed leaf forest", (160,180,50)), 92: ("Closed mixed leaf forest", (120,130,0)),
    120: ("Shrubland", (150,100,0)), 121: ("Evergreen shrubland", (150,75,0)),
    122: ("Deciduous shrubland", (150,100,0)), 130: ("Grassland", (255,180,50)),
    140: ("Lichens and mosses", (255,220,210)), 150: ("Sparse vegetation", (255,235,175)),
    152: ("Sparse shrubland", (255,210,120)), 153: ("Sparse herbaceous", (255,235,175)),
    181: ("Swamp", (0,168,132)), 182: ("Marsh", (115,255,223)),
    183: ("Flooded flat", (158,187,215)), 184: ("Saline", (130,130,130)),
    185: ("Mangrove", (245,122,182)), 186: ("Salt marsh", (102,205,171)),
    187: ("Tidal flat", (68,79,137)), 190: ("Impervious surfaces", (195,20,0)),
    200: ("Bare areas", (255,245,215)), 201: ("Consolidated bare areas", (220,220,220)),
    202: ("Unconsolidated bare areas", (255,245,215)), 210: ("Water body", (0,70,200)),
    220: ("Permanent ice and snow", (255,255,255)), 0: ("Filled / nodata", (255,255,255)),
}

SECTOR_NAMES = ['N','NE','E','SE','S','SW','W','NW']

# ---------- helpers ----------
def parse_tile_name(fname):
    base = os.path.basename(fname)
    m = re.search(r'([EW])\s*0*([0-9]{1,3}).*?([NS])\s*0*([0-9]{1,3})', base, re.IGNORECASE)
    if not m:
        raise ValueError(f"Cannot parse tile coords from filename: {base}")
    ew, lon_s, ns, lat_s = m.group(1).upper(), m.group(2), m.group(3).upper(), m.group(4)
    lon = int(lon_s); lat = int(lat_s)
    if ew == 'W': lon = -lon
    if ns == 'S': lat = -lat
    return lon, lat

def build_tile_index(tiles_dir):
    files = list(Path(tiles_dir).glob("**/*.tif"))
    idx = {}
    for f in files:
        try:
            lon, lat = parse_tile_name(str(f))
        except Exception:
            continue
        idx[(lon, lat)] = str(f)
    return idx

def deterministic_name(lon_c, lat_c):
    s = f"{lon_c:.6f},{lat_c:.6f}"
    h = hashlib.sha1(s.encode()).hexdigest()
    ci = int(h[0:6], 16) % len(colors)
    ai = int(h[6:12], 16) % len(animals)
    pi = int(h[12:18], 16) % len(places)
    return f"{colors[ci]} {animals[ai]} {places[pi]}"

def make_sector_masks(side):
    cx = (side - 1) / 2.0
    cy = (side - 1) / 2.0
    ys = np.arange(side); xs = np.arange(side)
    X, Y = np.meshgrid(xs, ys)
    dx = X - cx
    dy = cy - Y
    angles = np.degrees(np.arctan2(dy, dx))
    angle_from_north = (90 - angles) % 360
    sector = np.floor((angle_from_north + 22.5) / 45.0).astype(int) % 8
    return sector

def top_k_from_array(arr, k=3, ignore_values=None):
    if ignore_values is None: ignore_values = []
    if arr.size == 0:
        return []
    arr_flat = arr.ravel()
    if len(ignore_values)>0:
        arr_flat = arr_flat[~np.isin(arr_flat, ignore_values)]
    if arr_flat.size == 0:
        return []
    u, c = np.unique(arr_flat, return_counts=True)
    idx = np.argsort(-c)
    res = [(int(u[i]), int(c[i])) for i in idx[:k]]
    return res

def pixel_to_lonlat(lon_min, lat_max, px_deg, col, row):
    lon = lon_min + (col * px_deg)
    lat = lat_max - (row * px_deg)
    return lon, lat

def letters_from_index(idx):
    if idx < 0:
        idx = 0
    max_vals = 26 * 26
    n = idx % max_vals
    first = n // 26
    second = n % 26
    return chr(ord('A') + first) + chr(ord('A') + second)

# ---------- checkpoint utilities ----------
def load_checkpoint(csv_path):
    if not os.path.exists(csv_path):
        return {}
    d = {}
    with open(csv_path, newline='', encoding='utf8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            d[row['key']] = row
    return d

def append_checkpoint_row(csv_path, header, rowdict):
    write_header = not os.path.exists(csv_path)
    with open(csv_path, 'a', newline='', encoding='utf8') as f:
        writer = csv.DictWriter(f, fieldnames=header)
        if write_header:
            writer.writeheader()
        writer.writerow(rowdict)
        f.flush()
        os.fsync(f.fileno())

def update_checkpoint_row(csv_path, key, new_row):
    rows = []
    header = None
    if os.path.exists(csv_path):
        with open(csv_path, newline='', encoding='utf8') as f:
            reader = csv.DictReader(f)
            header = reader.fieldnames
            for r in reader:
                if r['key'] == key:
                    rows.append(new_row)
                else:
                    rows.append(r)
    else:
        header = list(new_row.keys())
        rows = [new_row]
    with open(csv_path, 'w', newline='', encoding='utf8') as f:
        writer = csv.DictWriter(f, fieldnames=header)
        writer.writeheader()
        for r in rows:
            writer.writerow(r)
        f.flush(); os.fsync(f.fileno())

# ---------- tile label (5° bins) ----------
def floor_to_5(x):
    # renvoie le multiple de 5 inférieur (ex: 45.7 -> 45 ; -117.3 -> -120)
    return math.floor(x / 5.0) * 5

def tile_label_from_lonlat_5deg(lon, lat):
    lon5 = int(floor_to_5(lon))
    lat5 = int(floor_to_5(lat))
    if lon5 >= 0:
        lon_label = f"E{lon5}"
    else:
        lon_label = f"W{abs(lon5)}"
    if lat5 >= 0:
        lat_label = f"N{lat5}"
    else:
        lat_label = f"S{abs(lat5)}"
    return f"{lon_label}{lat_label}"  # ex: E45N50 or W45S5

# ---------- JSON per-tile writer (robuste pour Windows) ----------
def append_base_to_tile_json(tile_label, base_obj):
    os.makedirs(JSON_OUT_DIR, exist_ok=True)
    path = os.path.join(JSON_OUT_DIR, f"{tile_label}.json")

    # Charger l'existant (si présent)
    data = []
    if os.path.exists(path):
        try:
            with open(path, 'r', encoding='utf8') as f:
                data = json.load(f)
                if not isinstance(data, list):
                    data = []
        except Exception:
            # corrupted or locked for reading: try to continue with empty list
            data = []

    data.append(base_obj)

    # écrire dans un tmp dans le même dossier, puis remplacer
    dirpath = os.path.dirname(path) or "."
    fd, tmpname = tempfile.mkstemp(prefix=".__tmp_json_", suffix=".json.tmp", dir=dirpath)
    os.close(fd)
    try:
        with open(tmpname, 'w', encoding='utf8') as tf:
            json.dump(data, tf, ensure_ascii=False, indent=2)
            tf.flush(); os.fsync(tf.fileno())

        # tenter os.replace avec retries
        last_err = None
        for attempt in range(REPLACE_RETRIES):
            try:
                os.replace(tmpname, path)
                return path
            except PermissionError as e:
                last_err = e
                time.sleep(REPLACE_SLEEP)
            except OSError as e:
                last_err = e
                time.sleep(REPLACE_SLEEP)

        # fallback : essayer d'écrire directement sur le fichier cible (moins atomique)
        try:
            with open(path, 'w', encoding='utf8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
                f.flush(); os.fsync(f.fileno())
            # cleanup tmp if still present
            if os.path.exists(tmpname):
                try: os.remove(tmpname)
                except: pass
            return path
        except Exception as e:
            # si tout échoue, supprimer tmp et remonter l'erreur originale
            if os.path.exists(tmpname):
                try: os.remove(tmpname)
                except: pass
            raise last_err if last_err is not None else e
    finally:
        # assure la suppression du tmp résiduel si quelque chose s'est mal passé
        if os.path.exists(tmpname):
            try: os.remove(tmpname)
            except: pass

# ---------- pipeline principal ----------
def process_all_tiles(tiles_dir):
    tile_index = build_tile_index(tiles_dir)
    if not tile_index:
        raise RuntimeError("Aucune tuile trouvée dans " + tiles_dir)

    sample_tile = next(iter(tile_index.values()))
    with rasterio.open(sample_tile) as s0:
        sample_px_x = abs(s0.transform.a)
        sample_px_y = abs(s0.transform.e)
        sample_px_deg = (sample_px_x + sample_px_y) / 2.0
        sample_band_count = s0.count

    global BAND_TO_PREFER
    if BAND_TO_PREFER is None:
        BAND_TO_PREFER = sample_band_count

    total_bases = 0
    tile_base_counts = {}
    for (lon_origin, lat_origin), path in tile_index.items():
        with rasterio.open(path) as ds:
            n_cols = math.ceil(ds.width / PIX_SIDE)
            n_rows = math.ceil(ds.height / PIX_SIDE)
            count = n_cols * n_rows
            tile_base_counts[(lon_origin, lat_origin)] = (n_rows, n_cols)
            total_bases += count

    checkpoint = load_checkpoint(CHECKPOINT_CSV)
    sector_mask = make_sector_masks(PIX_SIDE)
    summary = {"meta": {"ocean_code": OCEAN_CODE}, "bases_count": 0}
    ds_cache = {}
    def open_ds(key):
        if key not in ds_cache:
            p = tile_index.get(key)
            if p is None:
                return None
            ds_cache[key] = rasterio.open(p)
        return ds_cache[key]

    tile_keys = sorted(tile_index.keys(), key=lambda x: (-x[1], x[0]))  # lat desc, lon asc
    pbar = tqdm(total=total_bases, desc="total bases")

    for tile_key in tile_keys:
        lon_origin, lat_origin = tile_key
        ds = open_ds(tile_key)
        if ds is None:
            n_rows, n_cols = tile_base_counts.get(tile_key, (0,0))
            pbar.update(n_rows * n_cols)
            continue
        w = ds.width; h = ds.height
        transform = ds.transform
        lon_min = transform.c
        lat_max = transform.f
        px_x = abs(transform.a); px_y = abs(transform.e)
        px_deg_tile = (px_x + px_y) / 2.0
        band_count = ds.count
        ds_nodata = ds.nodatavals[0] if isinstance(ds.nodatavals, (list,tuple)) else ds.nodata
        nodata_vals = set(NODATA_DEFAULTS)
        if ds_nodata is not None:
            try: nodata_vals.add(int(ds_nodata))
            except: pass

        n_rows, n_cols = tile_base_counts[(lon_origin, lat_origin)]
        base_index = 0

        for r in range(n_rows):
            for c in range(n_cols):
                col_off = c * PIX_SIDE
                row_off = r * PIX_SIDE
                tile_lon_label = f"E{int(lon_origin)}" if lon_origin >= 0 else f"W{abs(int(lon_origin))}"
                tile_lat_label = f"N{int(lat_origin)}" if lat_origin >= 0 else f"S{abs(int(lat_origin))}"
                idx_letters = letters_from_index(base_index)
                key = f"{tile_lon_label}_{tile_lat_label}_{idx_letters}"

                cp = checkpoint.get(key)
                if cp and cp.get('status') == 'done':
                    pbar.update(1)
                    base_index += 1
                    continue

                row_cp = {
                    "key": key,
                    "tile_lon": lon_origin,
                    "tile_lat": lat_origin,
                    "col_off": col_off,
                    "row_off": row_off,
                    "center_lat": "",
                    "center_lon": "",
                    "name": "",
                    "status": "processing",
                    "ndjson_line_offset": ""
                }
                if not cp:
                    header = list(row_cp.keys())
                    append_checkpoint_row(CHECKPOINT_CSV, header, row_cp)
                else:
                    cp.update(row_cp)
                    update_checkpoint_row(CHECKPOINT_CSV, key, cp)

                # lecture / assemblage des bandes
                bands_needed = np.zeros((band_count, PIX_SIDE, PIX_SIDE), dtype=np.int32)
                inside_col0 = max(0, col_off)
                inside_row0 = max(0, row_off)
                inside_col1 = min(w, col_off + PIX_SIDE)
                inside_row1 = min(h, row_off + PIX_SIDE)
                if inside_col1 > inside_col0 and inside_row1 > inside_row0:
                    win = Window(inside_col0, inside_row0, inside_col1 - inside_col0, inside_row1 - inside_row0)
                    try:
                        arr_all = ds.read(list(range(1, band_count+1)), window=win)
                    except:
                        arr_all = ds.read(list(range(1, band_count+1)))
                    dst_col0 = inside_col0 - col_off
                    dst_row0 = inside_row0 - row_off
                    bands_needed[:, dst_row0:dst_row0 + arr_all.shape[1], dst_col0:dst_col0 + arr_all.shape[2]] = arr_all.astype(np.int32)

                neigh_offsets = [(0,0),(5,0),(-5,0),(0,5),(0,-5),(5,5),(5,-5),(-5,5),(-5,-5)]
                for dlon_deg, dlat_deg in neigh_offsets:
                    if dlon_deg == 0 and dlat_deg == 0:
                        continue
                    neighbor_origin = (lon_origin + dlon_deg, lat_origin + dlat_deg)
                    if neighbor_origin not in tile_index:
                        continue
                    nds = open_ds(neighbor_origin)
                    if nds is None:
                        continue
                    nw = nds.width; nh = nds.height
                    delta_lon = neighbor_origin[0] - lon_origin
                    delta_lat = neighbor_origin[1] - lat_origin
                    col_shift = int(round(delta_lon / px_deg_tile))
                    row_shift = int(round(-delta_lat / px_deg_tile))
                    neigh_col_origin = -col_shift
                    neigh_row_origin = -row_shift
                    a0 = col_off; a1 = col_off + PIX_SIDE
                    b0 = neigh_col_origin; b1 = neigh_col_origin + nw
                    inter_c0 = max(a0, b0); inter_c1 = min(a1, b1)
                    a0r = row_off; a1r = row_off + PIX_SIDE
                    b0r = neigh_row_origin; b1r = neigh_row_origin + nh
                    inter_r0 = max(a0r, b0r); inter_r1 = min(a1r, b1r)
                    if inter_c1 > inter_c0 and inter_r1 > inter_r0:
                        neigh_col_off = inter_c0 - neigh_col_origin
                        neigh_row_off = inter_r0 - neigh_row_origin
                        neigh_w = inter_c1 - inter_c0
                        neigh_h = inter_r1 - inter_r0
                        try:
                            nar = nds.read(list(range(1, band_count+1)), window=Window(neigh_col_off, neigh_row_off, neigh_w, neigh_h))
                        except:
                            nar = nds.read(list(range(1, band_count+1)))
                        dst_col0 = inter_c0 - col_off
                        dst_row0 = inter_r0 - row_off
                        bands_needed[:, dst_row0:dst_row0 + nar.shape[1], dst_col0:dst_col0 + nar.shape[2]] = nar.astype(np.int32)

                final = np.zeros((PIX_SIDE, PIX_SIDE), dtype=np.int32)
                nodata_vals_combined = set(NODATA_DEFAULTS)
                if ds_nodata is not None:
                    try:
                        nodata_vals_combined.add(int(ds_nodata))
                    except:
                        pass
                for b_idx in reversed(range(band_count)):
                    band_arr = bands_needed[b_idx]
                    mask_valid = (~np.isin(band_arr, list(nodata_vals_combined))) & (final == 0)
                    final[mask_valid] = band_arr[mask_valid]
                final[final == 0] = OCEAN_CODE

                # coordonnées
                center_col = col_off + (PIX_SIDE // 2)
                center_row = row_off + (PIX_SIDE // 2)
                center_lon, center_lat = pixel_to_lonlat(lon_min, lat_max, px_deg_tile, center_col, center_row)

                name = deterministic_name(center_lon, center_lat)

                # secteurs : uniquement codes (TOPK), pas de counts ni labels
                sectors_info = {}
                for si, sname in enumerate(SECTOR_NAMES):
                    mask = (sector_mask == si)
                    arr_sector = final[mask]
                    topk = top_k_from_array(arr_sector, k=TOPK, ignore_values=[])
                    codes_only = [int(code) for code, cnt in topk]
                    sectors_info[sname] = codes_only

                entry_minimal = {
                    "key": key,
                    "name": name,
                    "center": [center_lat, center_lon],
                    "sectors": sectors_info
                }

                # determine tile file (5° bins)
                center_tile_label = tile_label_from_lonlat_5deg(center_lon, center_lat)  # ex: E45N50
                written_path = append_base_to_tile_json(center_tile_label, entry_minimal)

                # update checkpoint
                cp_done = {
                    "key": key,
                    "tile_lon": lon_origin,
                    "tile_lat": lat_origin,
                    "col_off": col_off,
                    "row_off": row_off,
                    "center_lat": center_lat,
                    "center_lon": center_lon,
                    "name": name,
                    "status": "done",
                    "ndjson_line_offset": os.path.basename(written_path)
                }
                update_checkpoint_row(CHECKPOINT_CSV, key, cp_done)
                checkpoint[key] = cp_done

                summary["bases_count"] = summary.get("bases_count", 0) + 1
                if summary["bases_count"] % 100 == 0:
                    with open(OUTPUT_SUMMARY, "w", encoding="utf8") as sf:
                        json.dump(summary, sf, ensure_ascii=False, indent=2)

                base_index += 1
                pbar.update(1)

    for ds in ds_cache.values():
        try: ds.close()
        except: pass
    with open(OUTPUT_SUMMARY, "w", encoding="utf8") as sf:
        json.dump(summary, sf, ensure_ascii=False, indent=2)
    pbar.close()
    print("Traitement terminé. JSON par tuile dans:", JSON_OUT_DIR, "Checkpoint:", CHECKPOINT_CSV, "Summary:", OUTPUT_SUMMARY)


if __name__ == "__main__":
    process_all_tiles(TILES_DIR)