"""Summarise a matrix_capture.py run: mix, fight statistics, coverage maps and sample sheets.

usage: matrix_report.py DATA_DIR OUT_DIR

Writes OUT_DIR/report.json, coverage-<map>.png (top-down: grey = walkable
blocks from the plan, colour = player frames per 2 m cell, with eight yaw
sectors per 10 m block counted), and sheet-<map>.jpg (graybox | cartoon | war |
punk at mid-clip for a spread of team sizes).
"""
import collections
import json
import math
import random
import sys
from pathlib import Path

import cv2
import numpy as np

LOOKS = ('original', 'cartoon', 'war', 'punk')


def frame_at(video, index):
    cap = cv2.VideoCapture(str(video))
    cap.set(cv2.CAP_PROP_POS_FRAMES, index)
    ok, image = cap.read()
    cap.release()
    return image if ok else np.zeros((544, 960, 3), np.uint8)


def main():
    data, out = Path(sys.argv[1]), Path(sys.argv[2])
    out.mkdir(parents=True, exist_ok=True)
    jobs = [json.loads(l) for l in (data / 'jobs.jsonl').read_text().splitlines() if l.strip()]
    report = {}
    for map_id in ('paris', 'tsinghua'):
        metas = [json.loads((c / 'meta.json').read_text()) for c in sorted((data / map_id).glob('clip_*')) if (c / 'meta.json').exists()]
        fights = [m for m in metas if m['teams'][1]]
        positions, sectors = [], collections.defaultdict(set)
        for m in metas:
            for line in (data / map_id / f"clip_{m['clip']:06d}" / 'frames.jsonl').read_text().splitlines()[::6]:
                row = json.loads(line)
                p, yaw = row['state']['player'], row['state']['yaw']
                positions.append((p['x'], p['z']))
                sectors[(math.floor(p['x'] / 10), math.floor(p['z'] / 10))].add(int(((yaw % (2 * math.pi)) / (math.pi / 4))) % 8)
        planned = {(math.floor(j['setup']['player']['x'] / 10), math.floor(j['setup']['player']['z'] / 10)) for j in jobs if j['map'] == map_id}
        report[map_id] = dict(
            clips=len(metas), planned=sum(j['map'] == map_id for j in jobs),
            teams=dict(sorted(collections.Counter(f"{m['teams'][0]}v{m['teams'][1]}" for m in metas).items())),
            distance=dict(collections.Counter(m['distance'] for m in fights)), side=dict(collections.Counter(m['side'] for m in fights)),
            landmark_clips=sum(m['anchor'] is not None for m in metas), max_drift=max((m['max_drift'] for m in metas), default=None),
            fight_bullet_frame_share=round(float(np.mean([m['bullet_frame_share'] for m in fights])), 3) if fights else None,
            fight_mean_alive_bots=round(float(np.mean([m['mean_alive_bots'] for m in fights])), 2) if fights else None,
            fight_player_shots=round(float(np.mean([m['shots'] for m in fights])), 1) if fights else None,
            fight_kills=round(float(np.mean([m['kills'] for m in fights])), 2) if fights else None,
            blocks_visited=len(sectors), blocks_planned=len(planned),
            yaw_sectors_per_block=round(float(np.mean([len(s) for s in sectors.values()])), 2) if sectors else None,
            yaw_histogram=[sum(k in s for s in sectors.values()) for k in range(8)],
            wall_s_mean=round(float(np.mean([m['wall_s'] for m in metas])), 1) if metas else None)
        # Top-down coverage map: 2 m cells, log count colour.
        xs, zs = [p[0] for p in positions], [p[1] for p in positions]
        if positions:
            x0, x1, z0, z1 = min(xs) - 4, max(xs) + 4, min(zs) - 4, max(zs) + 4
            cell = 2.0
            grid = np.zeros((int((z1 - z0) / cell) + 1, int((x1 - x0) / cell) + 1))
            for x, z in positions:
                grid[int((z - z0) / cell), int((x - x0) / cell)] += 1
            norm = np.log1p(grid) / max(1e-6, np.log1p(grid).max())
            image = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
            image[grid == 0] = (40, 40, 40)
            for bx, bz in planned:
                x, z = int((bx * 10 - x0) / cell), int((bz * 10 - z0) / cell)
                cv2.rectangle(image, (x, z), (x + 5, z + 5), (90, 90, 90), 1)
            image = cv2.resize(image, None, fx=8, fy=8, interpolation=cv2.INTER_NEAREST)
            cv2.imwrite(str(out / f'coverage-{map_id}.png'), image[::-1])
        # Sample sheet: one fight per team size, mid-clip.
        rng = random.Random(0)
        by_team = collections.defaultdict(list)
        for m in metas:
            by_team[tuple(m['teams'])].append(m)
        rows = []
        for team in sorted(by_team)[:9]:
            m = rng.choice(by_team[team])
            clip = data / map_id / f"clip_{m['clip']:06d}"
            tiles = [cv2.resize(frame_at(clip / f'{look}.mp4', 120), (384, 218), interpolation=cv2.INTER_AREA) for look in LOOKS]
            row = np.concatenate(tiles, 1)
            cv2.putText(row, f"{m['teams'][0]}v{m['teams'][1]} {m['distance']} {m['side']} clip {m['clip']}", (8, 22), 0, .6, (0, 0, 0), 4)
            cv2.putText(row, f"{m['teams'][0]}v{m['teams'][1]} {m['distance']} {m['side']} clip {m['clip']}", (8, 22), 0, .6, (255, 255, 255), 2)
            rows.append(row)
        if rows:
            cv2.imwrite(str(out / f'sheet-{map_id}.jpg'), np.concatenate(rows, 0), [cv2.IMWRITE_JPEG_QUALITY, 82])
    (out / 'report.json').write_text(json.dumps(report, indent=1))
    print(json.dumps(report, indent=1))


if __name__ == '__main__':
    main()