insar-management-system-v2/experiments/isce2_sbas_timeseries/scripts/materialize_lt1_stack_scenes.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import shelve
import shutil
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional


def choose_path(scene: Dict[str, Any], primary_key: str, fallback_key: str) -> str:
    primary = scene.get(primary_key)
    fallback = scene.get(fallback_key)
    for candidate in (primary, fallback):
        if candidate and Path(candidate).exists():
            return str(Path(candidate))
    raise FileNotFoundError(
        f"Neither {primary_key} nor {fallback_key} exists for scene {scene.get('date') or scene.get('imaging_date')}"
    )


def load_manifest(path: Path) -> Dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def remove_existing_shelve(stem: Path) -> None:
    for suffix in ("", ".db", ".dat", ".dir", ".bak"):
        candidate = Path(str(stem) + suffix)
        if candidate.exists():
            if candidate.is_dir():
                shutil.rmtree(candidate)
            else:
                candidate.unlink()


def shelve_stem_exists(stem: Path) -> bool:
    for suffix in ("", ".db", ".dat", ".dir", ".bak"):
        if Path(str(stem) + suffix).exists():
            return True
    return False


@dataclass
class SceneResult:
    date: str
    output_dir: str
    slc_path: str
    data_shelve: str
    status: str
    bytes_written: Optional[int]
    started_at_utc: str
    ended_at_utc: str


def materialize_one_scene(
    scene: Dict[str, Any],
    force: bool,
) -> SceneResult:
    import isce
    from isceobj.Sensor import createSensor

    date = str(scene["date"])
    output_dir = Path(scene["target_dir_wsl"])
    output_dir.mkdir(parents=True, exist_ok=True)

    slc_path = Path(scene["expected_slc_wsl"])
    slc_xml_path = Path(scene["expected_slc_xml_wsl"])
    data_shelve = Path(scene["expected_data_shelve_wsl"])

    tiff_path = choose_path(scene, "source_tiff_wsl", "source_tiff_windows")
    orbit_xml = choose_path(scene, "orbit_xml_wsl", "orbit_xml_windows")

    if force:
        for path in (slc_path, slc_xml_path, Path(str(slc_path) + ".vrt")):
            if path.exists():
                path.unlink()
        remove_existing_shelve(data_shelve)

    if slc_path.exists() and slc_xml_path.exists() and shelve_stem_exists(data_shelve):
        now = datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
        return SceneResult(
            date=date,
            output_dir=str(output_dir),
            slc_path=str(slc_path),
            data_shelve=str(data_shelve),
            status="skipped_existing",
            bytes_written=slc_path.stat().st_size,
            started_at_utc=now,
            ended_at_utc=now,
        )

    started_at = datetime.utcnow().replace(microsecond=0).isoformat() + "Z"

    sensor = createSensor("LUTAN1")
    sensor.configure()
    sensor.tiff = tiff_path
    sensor.orbitFile = orbit_xml
    sensor.output = str(slc_path)
    sensor.extractImage()
    sensor.extractDoppler()
    sensor.frame.getImage().renderHdr()

    remove_existing_shelve(data_shelve)
    with shelve.open(str(data_shelve)) as db:
        db["frame"] = sensor.frame

    ended_at = datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
    report = {
        "date": date,
        "source_tiff": tiff_path,
        "orbit_xml": orbit_xml,
        "output_slc": str(slc_path),
        "output_slc_xml": str(slc_xml_path),
        "data_shelve": str(data_shelve),
        "frame_lines": sensor.frame.getNumberOfLines(),
        "frame_samples": sensor.frame.getNumberOfSamples(),
        "started_at_utc": started_at,
        "ended_at_utc": ended_at,
    }
    (output_dir / "materialization_report.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )

    return SceneResult(
        date=date,
        output_dir=str(output_dir),
        slc_path=str(slc_path),
        data_shelve=str(data_shelve),
        status="materialized",
        bytes_written=slc_path.stat().st_size if slc_path.exists() else None,
        started_at_utc=started_at,
        ended_at_utc=ended_at,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Materialize LT-1 stack acquisitions into stripmapStack-ready SLC/date directories."
    )
    parser.add_argument(
        "--stack-manifest",
        required=True,
        help="Path to stack_input_manifest.json generated by build_lt1_stack_prep.py",
    )
    parser.add_argument(
        "--dates",
        nargs="+",
        default=None,
        help="Optional subset of acquisition dates to materialize, for example 20250510 20250705",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Overwrite existing .slc/.xml/data outputs for the selected dates.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    manifest_path = Path(args.stack_manifest)
    if not manifest_path.exists():
        raise FileNotFoundError(f"Stack manifest not found: {manifest_path}")

    manifest = load_manifest(manifest_path)
    scenes = list(manifest.get("scenes", []))
    if not scenes:
        raise ValueError(f"No scenes found in stack manifest: {manifest_path}")

    selected_dates = set(args.dates or [])
    if selected_dates:
        scenes = [scene for scene in scenes if str(scene["date"]) in selected_dates]
        if not scenes:
            raise ValueError(f"No matching dates found in manifest for selection: {sorted(selected_dates)}")

    results: List[SceneResult] = []
    for scene in scenes:
        print(f"Materializing {scene['date']} -> {scene['target_dir_wsl']}")
        result = materialize_one_scene(scene, force=args.force)
        results.append(result)
        print(f"  status={result.status} slc={result.slc_path}")

    report = {
        "generated_at_utc": datetime.utcnow().replace(microsecond=0).isoformat() + "Z",
        "stack_manifest": str(manifest_path),
        "results": [result.__dict__ for result in results],
    }

    report_path = manifest_path.parent / "materialization_summary.json"
    report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"Summary: {report_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())