"""Read-only explanations layered over the unmodified B engine."""
import hashlib
from itertools import combinations
import json
from math import dist
from pathlib import Path
import sys

sys.path.append('/baseline_helpers')
from features import Settings, bodies, match_cost
from contract import observation, edge_state


def verify_baseline():
    expected = json.loads(Path('/baseline_hashes.json').read_text())
    for name, digest in expected.items():
        assert hashlib.sha256((Path('/baseline_helpers') / name).read_bytes()).hexdigest() == digest
    return expected


def load_rows(clip):
    previous = Path('/previous') / clip / 'data'
    check = json.loads((previous / 'cache_validation.json').read_text())
    raw = Path('/cache/inputs/pose.json').read_bytes()
    assert check['status'] == 'PASS' and hashlib.sha256(raw).hexdigest() == check['cache_sha256']
    all_rows = json.loads(raw)
    selected = [all_rows[i] for i in check['selected_samples']]
    chosen = {r['sample'] for r in selected}
    assert len(chosen) == len(selected)
    start = selected[0]['seconds'] - 3.
    rows = [r for r in all_rows if start - 1e-8 <= r['seconds'] <= selected[-1]['seconds'] + 1e-8]
    assert all(b['seconds'] > a['seconds'] and b['sample'] > a['sample'] for a, b in zip(rows, rows[1:]))
    return rows, chosen, check, previous


def strict_people(row, palette):
    people = observation(row, palette)
    for p in people:
        p['joint_states'] = ['BELOW' if s == 'LOW' else s for s in p['joint_states']]
        from contract import EDGES
        p['edge_states'] = [edge_state(p['joint_states'][a], p['joint_states'][b]) for a, b in EDGES]
        p['detail_slot'] = None
        p['fight_display'] = False
    return [p for p in people if p['visible']]


def hull(points):
    pts = sorted(set(tuple(p) for p in points))
    if len(pts) <= 1:
        return pts
    def cross(o, a, b):
        return (a[0]-o[0])*(b[1]-o[1]) - (a[1]-o[1])*(b[0]-o[0])
    lower, upper = [], []
    for p in pts:
        while len(lower) >= 2 and cross(lower[-2], lower[-1], p) <= 0:
            lower.pop()
        lower.append(p)
    for p in reversed(pts):
        while len(upper) >= 2 and cross(upper[-2], upper[-1], p) <= 0:
            upper.pop()
        upper.append(p)
    return lower[:-1] + upper[:-1]


def candidate_pairs(bs, cfg):
    return [p for p in combinations(bs, 2)
        if dist(p[0]['center'], p[1]['center']) / ((p[0]['scale']+p[1]['scale'])/2) <= cfg.pair_distance_ratio][:cfg.max_pairs]


def spatial_groups(bs, pairs, width, height):
    by_order = {b['order']: b for b in bs}
    graph = {}
    edges = []
    for a, b in pairs:
        x, y = a['order'], b['order']
        graph.setdefault(x, set()).add(y)
        graph.setdefault(y, set()).add(x)
        edges.append([x, y])
    groups, seen = [], set()
    for start in sorted(graph):
        if start in seen:
            continue
        stack, members = [start], set()
        while stack:
            x = stack.pop()
            if x in members:
                continue
            members.add(x)
            stack.extend(graph[x] - members)
        seen.update(members)
        points = []
        for x in members:
            a, b, c, d = by_order[x]['bbox']
            points += [(max(0,min(width-1,xx)), max(0,min(height-1,yy))) for xx, yy in ((a,b),(c,b),(c,d),(a,d))]
        groups.append(dict(group='G'+str(len(groups)+1), orders=sorted(members),
            edges=[e for e in edges if e[0] in members], polygon=hull(points), frame_local=True))
    return groups


def valid_ids(ids):
    return len(ids) == 2 and all(i is not None for i in ids) and ids[0] != ids[1]


def association(pair, before, now, cfg):
    token, ids = pair['token'], pair['ids']
    previous = before.get(token)
    same_t = [dict(token=k, **v) for k, v in before.items()
        if valid_ids(ids) and valid_ids([b['tid'] for b in v['bodies']])
        and set(ids) == {b['tid'] for b in v['bodies']}
        and 0 < now-v['time'] <= cfg.max_step_gap+1e-8]
    out = dict(kind='NEW_PAIR_OR_ELIGIBILITY_CHANGE', difference=False, token=token,
        current_ids=ids, orders=pair['orders'], previous_token=None, previous_ids=None,
        previous_time=None, elapsed=None, previous_bodies=None, costs=None,
        same_t_previous=[s['token'] for s in same_t])
    if previous is None:
        if same_t:
            previous = max(same_t, key=lambda s: s['time'])
            out.update(kind='B_RESET_SAME_T_PAIR', difference=True, previous_token=previous['token'])
        elif not valid_ids(ids):
            out['kind'] = 'UNCOMPARABLE_T_MISSING'
        else:
            older = [dict(token=k, **v) for k,v in before.items()
                if valid_ids([b['tid'] for b in v['bodies']]) and set(ids)=={b['tid'] for b in v['bodies']}]
            if older:
                previous=max(older,key=lambda s:s['time'])
                out.update(kind='OBSERVATION_GAP',previous_token=previous['token'])
    else:
        elapsed = now-previous['time']
        old_ids = [b['tid'] for b in previous['bodies']]
        out['previous_token'] = token
        if elapsed > cfg.max_step_gap+1e-8:
            out['kind'] = 'OBSERVATION_GAP'
        elif not valid_ids(ids) or not valid_ids(old_ids):
            out['kind'] = 'UNCOMPARABLE_T_MISSING'
        elif ids != old_ids:
            out.update(kind='T_SWAP_B_CONTINUES' if set(ids)==set(old_ids) else 'T_CHANGE_B_CONTINUES', difference=True)
        else:
            out['kind'] = 'AGREEMENT'
    if previous:
        out.update(previous_ids=[b['tid'] for b in previous['bodies']],
            previous_time=previous['time'], elapsed=now-previous['time'], previous_bodies=previous['bodies'])
    return out


def arm_evidence(arm, rules, cfg):
    checks = dict(valid=arm['valid'], motion_strike=arm['relative_motion'] >= cfg.strike_motion,
        toward=arm['toward'] >= cfg.strike_toward, approach=arm['approach'] >= cfg.strike_approach,
        hand_strike=arm['hand_distance'] <= cfg.strike_hand_distance,
        hand_contact=arm['hand_distance'] <= cfg.contact_hand_distance,
        motion_contact=arm['relative_motion'] >= cfg.contact_motion)
    partial = arm['valid'] and (checks['motion_strike'] or checks['motion_contact'] or checks['hand_contact'])
    instant = arm['strike'] or arm['dynamic']
    accumulated = ('STRIKE' in rules and arm['strike']) or ('PUSH_GRAPPLE' in rules and arm['dynamic'])
    return dict(actor=arm['actor'], side=arm['side'], level=3 if accumulated else 2 if instant else 1 if partial else 0,
        partial=bool(partial), instant=bool(instant), accumulated=bool(accumulated), checks=checks,
        values=arm, failed_strike=[k for k in ('valid','motion_strike','toward','approach','hand_strike') if not checks[k]])


class Observer:
    def __init__(self, cfg):
        self.cfg=cfg
        self.last_support={}
        self.last_support_time=None

    def capture(self, engine):
        return {k: dict(time=s.last_seen, bodies=s.last_bodies) for k,s in engine.states.items() if s.last_seen is not None}

    def explain(self, row, prepared, result, before):
        cfg=self.cfg
        bs=prepared[0]
        by_order={b['order']:b for b in bs}
        candidates=candidate_pairs(bs,cfg)
        assert len(candidates)==result['candidate_pairs']
        pair_info=[]
        for p in result['pairs']:
            current=[by_order[o] for o in p['orders']]
            a=association(p,before,row['seconds'],cfg)
            if a['previous_bodies']:
                a['costs'] = dict(actual_orientation=match_cost(a['previous_bodies'],current,cfg),
                    reversed_orientation=match_cost(a['previous_bodies'],current[::-1],cfg))
            evidence=[arm_evidence(x,p['rules'],cfg) for x in (p['features']['arms'] if p['features'] else [])]
            if p['features']:
                for e in evidence:
                    e['checks']['pair_change']=p['features']['pair_relative_change']>=cfg.pair_relative_change
                    e['failed_dynamic']=[k for k in ('valid','hand_contact','motion_contact','pair_change') if not e['checks'][k]]
            pair_info.append(dict(token=p['token'], orders=p['orders'], ids=p['ids'], refs=p['refs'], bodies=current,
                rules=p['rules'], hit=p['hit'], evidence=evidence, level=max((e['level'] for e in evidence),default=0),
                reason=p['reason'], association=a, input_start=p['input_start'], input_end=p['input_end'],
                pulses=p['pulses'], contact_seconds=p['contact_seconds'], dynamic_seconds=p['dynamic_seconds'],
                dynamic_samples=p['dynamic_samples'], close_seconds=p['close_seconds'], features=p['features']))
        if result['current_rules']:
            self.last_support={p['token']:row['seconds'] for p in pair_info if p['hit']}
            self.last_support_time=row['seconds']
        active={p['token'] for p in pair_info if p['hit']}
        held=set()
        if result['episode']['display'] and not active and self.last_support_time is not None:
            for p in pair_info:
                a=p['association']
                if p['token'] in self.last_support and a['previous_token']==p['token'] and a['elapsed'] is not None and a['elapsed']<=cfg.max_step_gap+1e-8:
                    held.add(p['token'])
            self.last_support={k:v for k,v in self.last_support.items() if k in held}
        elif not result['episode']['display']:
            self.last_support={}
        current_sets={tuple(sorted(p['ids'])) for p in pair_info if valid_ids(p['ids'])}
        ends=[]
        valid_t={b['tid'] for b in bs if b['tid'] is not None}
        for token,old in before.items():
            old_ids=[b['tid'] for b in old['bodies']]
            if valid_ids(old_ids) and tuple(sorted(old_ids)) not in current_sets:
                reason='PROXIMITY_OR_CANDIDATE_EXIT' if set(old_ids)<=valid_t else 'INPUT_UNAVAILABLE_OR_EXIT'
                ends.append(dict(token=token,ids=old_ids,reason=reason))
        for p in pair_info:
            p['highlight']='CURRENT_RULE' if p['token'] in active else 'DISPLAY_HOLD_CONFIRMED' if p['token'] in held else None
        return dict(sample=row['sample'],source_frame=row['source_frame'],seconds=row['seconds'],
            groups=spatial_groups(bs,candidates,row['width'],row['height']), bodies=bs,
            candidate_edges=[[a['order'],b['order']] for a,b in candidates], pairs=pair_info,
            current_support=sorted(active), held_support=sorted(held), departed_pairs=ends,
            unknown_hold=bool(result['episode']['display'] and not active and not held),
            episode=result['episode'], events=result['events'], current_rules=result['current_rules'],
            event_count=result['event_count'],rejected=result['rejected'],displayed=result['displayed'])