"""Отрисовка схемы мозга мухи по кадрам бэга — проверка и материал для видео.

    python tools/render_brain.py --bag data/for_hackathon/doubleT_obstacle --frames 8
    python tools/render_brain.py --bag ... --video docs/figures/brain.mp4 --limit 200
"""
from __future__ import annotations

import argparse

import numpy as np

import _bootstrap as B  # noqa: F401
from flyguard.bag import Bag
from flyguard.brain_view import BrainView
from flyguard.mushroom_body import MushroomBody
from flyguard.pipeline import FlyGuard, Params


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", required=True)
    ap.add_argument("--memory")
    ap.add_argument("--limit", type=int, default=120)
    ap.add_argument("--frames", type=int, default=6, help="сколько кадров сохранить в PNG")
    ap.add_argument("--video", default=None)
    ap.add_argument("--style", default="hybrid",
                    choices=("scheme", "cloud", "hybrid"),
                    help="схема | облако нейронов | панели + облако")
    args = ap.parse_args()

    if args.style == "scheme":
        view = BrainView()
    elif args.style == "cloud":
        from flyguard.brain_atlas import NeuronCloud
        view = NeuronCloud()
    else:
        from flyguard.brain_hybrid import BrainHybrid
        view = BrainHybrid()
    if not view.enabled:
        raise SystemExit("нужен OpenCV, а для облака — ещё атлас нейронов "
                         "(python tools/build_brain_atlas.py)")
    import cv2

    memory = MushroomBody.load(args.memory) if args.memory else None
    fg = FlyGuard(Params(enable_looming=True), memory=memory)
    bag = Bag(args.bag)

    B.FIGURES.mkdir(parents=True, exist_ok=True)
    writer = None
    saved = 0
    processed = 0
    # кадры калибровки решётки не доходят до отрисовки, поэтому шаг считается
    # от реально обрабатываемых, иначе при малом --limit не сохранится ничего
    usable = max(args.limit - Params().calib_frames, 1)
    every = max(1, usable // max(args.frames, 1))

    for k, (_, pc) in enumerate(bag.frames(stop=args.limit)):
        res = fg.process(pc, keep_debug=True)
        if res is None:
            continue
        processed += 1
        img = view.render(res)
        if img is None:
            continue
        if args.video:
            if writer is None:
                fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                writer = cv2.VideoWriter(args.video, fourcc, 10.0,
                                         (img.shape[1], img.shape[0]))
            writer.write(img)
        if processed % every == 0 and saved < args.frames:
            out = B.FIGURES / f"brain_{args.style}_{bag.path.name}_{k:04d}.png"
            cv2.imwrite(str(out), img)
            saved += 1
            print(f"  кадр {k:4d}: {out.name}  "
                  f"{'ПРЕПЯТСТВИЕ %.0f м' % res.decision.distance if res.decision.detected else 'путь свободен'}"
                  f"  {res.total_ms:.1f} мс")

    if writer is not None:
        writer.release()
        print("видео:", args.video)
    print(f"обработано кадров: {processed}, сохранено картинок: {saved}")


if __name__ == "__main__":
    main()