"""Per-frame structural agreement between the graybox and each styled look of paired matrix clips.

usage: matrix_qa.py DATA_DIR [DATA_DIR ...] [--clips 40]

For the same clips in every DATA_DIR (the first --clips per map, sorted), every
frame of every look is reduced to a blurred 160x90 edge-strength map (Sobel on
grey); the Pearson correlation between the graybox map and each styled map
measures whether the same geometry is on screen. A frame whose geometry was
dropped in one look scores far lower than its neighbours. Also reports the
share of the graybox's bottom third that is sky colour.
"""
import argparse
import json
from pathlib import Path

import cv2
import numpy as np

LOOKS = ('cartoon', 'war', 'punk')
SKY = {'paris': np.array([192, 168, 121]), 'tsinghua': np.array([196, 171, 124])}


def edges(frame):
    g = cv2.cvtColor(cv2.resize(frame, (160, 90), interpolation=cv2.INTER_AREA), cv2.COLOR_BGR2GRAY).astype(np.float32)
    m = cv2.magnitude(cv2.Sobel(g, cv2.CV_32F, 1, 0), cv2.Sobel(g, cv2.CV_32F, 0, 1))
    m = cv2.GaussianBlur(m, (0, 0), 2)
    return (m - m.mean()) / (m.std() + 1e-6)


def frames(path):
    cap = cv2.VideoCapture(str(path))
    out = []
    while True:
        ok, f = cap.read()
        if not ok:
            break
        out.append(f)
    return out


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('data', type=Path, nargs='+')
    parser.add_argument('--clips', type=int, default=40)
    args = parser.parse_args()
    report = {}
    for data in args.data:
        corr = {look: [] for look in LOOKS}
        sky = []
        for map_id in ('paris', 'tsinghua'):
            for clip in sorted((data / map_id).glob('clip_*'))[:args.clips]:
                gray = frames(clip / 'original.mp4')
                g_edges = [edges(f) for f in gray]
                for f in gray[::24]:
                    b = f[2 * f.shape[0] // 3:].astype(int)
                    sky.append(float((np.abs(b - SKY[map_id]).sum(-1) < 12).mean()))
                for look in LOOKS:
                    for ge, f in zip(g_edges, frames(clip / f'{look}.mp4')):
                        corr[look].append(float((ge * edges(f)).mean()))
        report[str(data)] = dict(
            frames=len(corr['war']), graybox_sky_bottom_third=round(float(np.mean(sky)), 3),
            **{look: dict(mean=round(float(np.mean(v)), 3), p05=round(float(np.percentile(v, 5)), 3),
                          below_0_3=round(float(np.mean(np.array(v) < .3)), 4)) for look, v in corr.items()})
    print(json.dumps(report, indent=1))


if __name__ == '__main__':
    main()