"""Full distribution statistics for a matrix_capture.py run (clip, frame and coverage level).

usage: matrix_stats.py DATA_DIR OUT_DIR

Reads every clip's meta.json, frames.jsonl and graybox cond frames. Writes
OUT_DIR/stats.json and per map: frames-<map>.png (player frames per 2 m cell),
sectors-<map>.png (yaw sectors seen per 10 m block, 0-8).

"In view" counts living characters whose direction from the camera lies inside
the horizontal field of view (80 degrees, 70 when aiming) and within 40 m; walls
are not traced, so it is an upper bound on what is actually visible.
"""
import collections
import json
import math
import sys
from pathlib import Path

import cv2
import numpy as np

SKY_TOL = 12


def hist(values, edges):
    counts = [0] * (len(edges) - 1)
    for v in values:
        for i in range(len(edges) - 1):
            if edges[i] <= v < edges[i + 1]:
                counts[i] += 1
                break
    return counts


def share(n, d):
    return round(n / d, 4) if d else None


def forward(q):
    x, y, z, w = q
    # Rotate (0, 0, -1) by quaternion (x, y, z, w).
    return (-(2 * (x * z + w * y)), -(2 * (y * z - w * x)), -(1 - 2 * (x * x + y * y)))


def main():
    data, out = Path(sys.argv[1]), Path(sys.argv[2])
    out.mkdir(parents=True, exist_ok=True)
    stats = {}
    for map_id in ('paris', 'tsinghua'):
        clips = sorted(c for c in (data / map_id).glob('clip_*') if (c / 'meta.json').exists())
        metas = [json.loads((c / 'meta.json').read_text()) for c in clips]
        c = collections.Counter
        clip = dict(
            clips=len(metas), frames_per_clip=metas[0]['frames'], fps=metas[0]['fps'],
            teams=dict(sorted(c(f"{m['teams'][0]}v{m['teams'][1]}" for m in metas).items())),
            characters=dict(sorted(c(m['teams'][0] + m['teams'][1] for m in metas).items())),
            distance=dict(c(m['distance'] for m in metas if m['teams'][1])),
            side=dict(c(m['side'] for m in metas if m['teams'][1])),
            pitch_tier=dict(sorted(c(str(m['pitch_bias']) for m in metas).items())),
            anchors=dict(c(m['anchor'] or 'block' for m in metas).most_common()),
            start_yaw_sector=hist([(m['setup']['player']['yaw'] % (2 * math.pi)) / (math.pi / 4) for m in metas], list(range(9))),
            teams_by_distance={d: dict(sorted(c(f"{m['teams'][0]}v{m['teams'][1]}" for m in metas if m['distance'] == d and m['teams'][1]).items()))
                               for d in ('close', 'mid', 'far')},
            max_drift=max(m['max_drift'] for m in metas),
            kills_per_clip=dict(sorted(c(m['kills'] for m in metas if m['teams'][1]).items())),
            player_shots_per_clip=hist([m['shots'] for m in metas if m['teams'][1]], [0, 1, 10, 20, 30, 40, 1000]),
            player_hit_rate=share(sum(m['hits'] for m in metas), sum(m['shots'] for m in metas)),
            wall_s_mean=round(float(np.mean([m['wall_s'] for m in metas])), 1))
        frames = collections.Counter()
        yaw_frames, pitch_frames, cam_dist, near_enemy, in_view_enemy, in_view_any = [], [], [], [], [], []
        cells, sectors, block_frames = collections.Counter(), collections.defaultdict(set), collections.Counter()
        refills, respawns, displacement = 0, 0, []
        sky_clip = []
        top = collections.Counter()
        for cdir, m in zip(clips, metas):
            fight = m['teams'][1] > 0
            prev, start, last = None, None, None
            for line in (cdir / 'frames.jsonl').read_text().splitlines():
                row = json.loads(line)
                s, inp, cam = row['state'], row['input'], row['camera']
                p = s['player']
                frames['total'] += 1
                frames['fight'] += fight
                held = set(inp['down'])
                frames['bullets'] += s['bullets'] > 0
                frames['fight_bullets'] += fight and s['bullets'] > 0
                frames['grenade_flying'] += s['grenades'] > 0
                frames['reloading'] += s['reloading'] > 0
                frames[f"weapon{s['weapon']}"] += 1
                frames['avatar_hidden'] += not cam['v']
                start = start or (p['x'], p['z'])
                last = (p['x'], p['z'])
                yaw_frames.append((s['yaw'] % (2 * math.pi)) / (math.pi / 4))
                pitch_frames.append(s['pitch'])
                if cam['s']:
                    cam_dist.append(cam['s']['distance'])
                cells[(math.floor(p['x'] / 2), math.floor(p['z'] / 2))] += 1
                block = (math.floor(p['x'] / 10), math.floor(p['z'] / 10))
                block_frames[block] += 1
                sectors[block].add(int((s['yaw'] % (2 * math.pi)) / (math.pi / 4)) % 8)
                if prev:
                    if p['health'] > prev['player']['health'] + 50:
                        refills += 1
                    respawns += sum(1 for a, b in zip(s['actors'], prev['actors']) if a['alive'] and not b['alive'])
                # Who is in front of the camera (three.js frame: z = -nav z).
                fx, fy, fz = forward(cam['q'])
                horiz = math.hypot(fx, fz) or 1
                half = math.radians(cam['fov'] * 960 / 544 / 2) if False else math.radians(40)
                cx, cz = cam['p'][0], cam['p'][2]
                seen_enemy = seen_any = 0
                dists = []
                for a in s['actors']:
                    if not a['alive']:
                        continue
                    ax, az = a['x'], -a['z']
                    dx, dz = ax - cx, az - cz
                    d = math.hypot(dx, dz)
                    if a['team'] == 1:
                        dists.append(math.hypot(a['x'] - p['x'], a['z'] - p['z']))
                    if d < 40 and d > .3:
                        cos = (dx * fx + dz * fz) / (d * horiz)
                        if cos > math.cos(half):
                            seen_any += 1
                            seen_enemy += a['team'] == 1
                if fight:
                    in_view_enemy.append(seen_enemy)
                    if dists:
                        near_enemy.append(min(dists))
                in_view_any.append(seen_any)
                for k in ('j', 'k', 'c', 'w', 'a', 's', 'd'):
                    frames[f'key_{k}'] += k in held or (prev is not None and k in prev['held'])
                for k in ('shift', ' ', 'f', 'r', '1', '2', '3'):
                    frames[f'tap_{k}'] += k in inp['down']
                prev = dict(row['state'], held=held | (prev['held'] - set(inp['up']) if prev else set()))
            displacement.append(math.hypot(last[0] - start[0], last[1] - start[1]))
            # Graybox ground check: share of the bottom third that is sky colour.
            shares = []
            for l in (0, 15, 30, 45, 59):
                im = cv2.imread(str(cdir / 'cond' / f'{l:06d}.png'))
                if im is None:
                    continue
                if len(top) < 2000:
                    top.update(map(tuple, im[0, ::40].tolist()))
                shares.append(im)
            sky_clip.append(shares)
        sky = np.array(top.most_common(1)[0][0])
        sky_share = []
        for images in sky_clip:
            vals = [float((np.abs(im[2 * im.shape[0] // 3:].astype(int) - sky).sum(-1) < SKY_TOL).mean()) for im in images]
            sky_share.append(float(np.mean(vals)) if vals else 0.0)
        n = frames['total']
        sector_counts = [len(v) for v in sectors.values()]
        bf = sorted(block_frames.values())
        frame = dict(
            frames=n, fight_frames=frames['fight'],
            yaw_sector_frame_share=[share(x, n) for x in hist(yaw_frames, list(range(9)))],
            pitch_hist=dict(zip(['<-0.3', '-0.3~-0.1', '-0.1~0.05', '0.05~0.2', '>=0.2'],
                                [share(x, n) for x in hist(pitch_frames, [-9, -.3, -.1, .05, .2, 9])])),
            camera_distance_hist=dict(zip(['<1.15 (avatar hidden)', '1.15~2.4', '2.4~4', '>=4'],
                                          [share(x, len(cam_dist)) for x in hist(cam_dist, [0, 1.15, 2.4, 4, 99])])),
            blocks_visited=len(sectors), block_frames_min_median_max=[bf[0], bf[len(bf) // 2], bf[-1]],
            blocks_by_yaw_sectors=dict(sorted(collections.Counter(sector_counts).items())),
            blocks_all_8_sectors=share(sum(x == 8 for x in sector_counts), len(sector_counts)),
            cells_2m_visited=len(cells),
            fight_frames_with_bullets=share(frames['fight_bullets'], frames['fight']),
            fight_enemies_in_view=dict(zip(['0', '1', '2', '3', '4+'], [share(x, len(in_view_enemy)) for x in hist(in_view_enemy, [0, 1, 2, 3, 4, 99])])),
            characters_in_view=dict(zip(['0', '1', '2', '3', '4+'], [share(x, len(in_view_any)) for x in hist(in_view_any, [0, 1, 2, 3, 4, 99])])),
            nearest_enemy_m=dict(zip(['<5', '5~10', '10~20', '20~35', '>=35'], [share(x, len(near_enemy)) for x in hist(near_enemy, [0, 5, 10, 20, 35, 999])])),
            fire_held=share(frames['key_j'], n), aim_held=share(frames['key_k'], n), crouch=share(frames['key_c'], n),
            move_w=share(frames['key_w'], n), strafe=share(frames['key_a'] + frames['key_d'], n), back_s=share(frames['key_s'], n),
            reloading=share(frames['reloading'], n), grenade_in_air=share(frames['grenade_flying'], n),
            weapon_share={'pistol': share(frames['weapon0'], n), 'rifle': share(frames['weapon1'], n), 'shotgun': share(frames['weapon2'], n)},
            avatar_hidden=share(frames['avatar_hidden'], n),
            events_per_clip={'dash': round(frames['tap_shift'] / len(metas), 2), 'jump': round(frames['tap_ '] / len(metas), 2),
                             'grenade': round(frames['tap_f'] / len(metas), 2), 'reload_key': round(frames['tap_r'] / len(metas), 2),
                             'weapon_switch': round((frames['tap_1'] + frames['tap_2'] + frames['tap_3']) / len(metas), 2),
                             'bot_respawn': round(respawns / len(metas), 2), 'player_refill': round(refills / len(metas), 2)},
            displacement_m=dict(zip(['<5', '5~15', '15~30', '>=30'], hist(displacement, [0, 5, 15, 30, 999]))),
            graybox_sky_bottom_third=dict(mean=round(float(np.mean(sky_share)), 3),
                                          clips=dict(zip(['<10%', '10~30%', '30~50%', '>=50%'], hist(sky_share, [0, .1, .3, .5, 1.01])))))
        stats[map_id] = dict(clip=clip, frame=frame)
        # Maps: frames per 2 m cell, and yaw sectors per 10 m block.
        xs = [k[0] for k in cells]
        zs = [k[1] for k in cells]
        grid = np.zeros((max(zs) - min(zs) + 1, max(xs) - min(xs) + 1))
        for (x, z), v in cells.items():
            grid[z - min(zs), x - min(xs)] = v
        norm = np.log1p(grid) / np.log1p(grid).max()
        image = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_VIRIDIS)
        image[grid == 0] = (245, 245, 245)
        cv2.imwrite(str(out / f'frames-{map_id}.png'), cv2.resize(image[::-1], None, fx=6, fy=6, interpolation=cv2.INTER_NEAREST))
        bx = [k[0] for k in sectors]
        bz = [k[1] for k in sectors]
        sec = np.full((max(bz) - min(bz) + 1, max(bx) - min(bx) + 1), -1)
        for (x, z), v in sectors.items():
            sec[z - min(bz), x - min(bx)] = len(v)
        image = cv2.applyColorMap((np.clip(sec, 0, 8) * 31).astype(np.uint8), cv2.COLORMAP_VIRIDIS)
        image[sec < 0] = (245, 245, 245)
        image = cv2.resize(image[::-1], None, fx=30, fy=30, interpolation=cv2.INTER_NEAREST)
        for (x, z), v in sectors.items():
            row = sec.shape[0] - 1 - (z - min(bz))
            cv2.putText(image, str(len(v)), ((x - min(bx)) * 30 + 9, row * 30 + 21), 0, .5, (255, 255, 255) if len(v) < 6 else (0, 0, 0), 1)
        cv2.imwrite(str(out / f'sectors-{map_id}.png'), image)
        print(map_id, 'done', flush=True)
    (out / 'stats.json').write_text(json.dumps(stats, indent=1, ensure_ascii=False))
    print(json.dumps(stats, ensure_ascii=False))


if __name__ == '__main__':
    main()