"""Dependency-free, partially observed navigation; no learned policy or MiniGrid run."""
import argparse
from collections import deque
from pathlib import Path
import hashlib
import json
import random
import statistics
import sys

ROOT = Path(__file__).resolve().parent
DELTAS = ((0, -1), (1, 0), (0, 1), (-1, 0))  # N,E,S,W; x,y
POLICIES = ('random', 'memory_reset', 'memory', 'oracle')


def dump(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2, ensure_ascii=False) + '\n')


def maze(seed, size=9):
    rng = random.Random(seed)
    grid = [['#'] * size for _ in range(size)]
    grid[1][1] = '.'
    stack = [(1, 1)]
    while stack:
        x, y = stack[-1]
        choices = [(dx, dy) for dx, dy in DELTAS
                   if 0 < x + 2*dx < size-1 and 0 < y + 2*dy < size-1
                   and grid[y+2*dy][x+2*dx] == '#']
        if not choices:
            stack.pop()
            continue
        dx, dy = rng.choice(choices)
        grid[y+dy][x+dx] = grid[y+2*dy][x+2*dx] = '.'
        stack.append((x+2*dx, y+2*dy))
    rows = [''.join(row) for row in grid]
    cells = {(x, y) for y, row in enumerate(rows) for x, v in enumerate(row) if v == '.'}
    routes = paths(cells, (1, 1))
    goals = sorted(p for p in cells if p[0] % 2 == 1 and p[1] % 2 == 1 and len(routes[p]) >= 8)
    goal = list(rng.choice(goals))
    return dict(seed=seed, rows=rows, start=[1, 1], goal=goal,
                map_sha256=hashlib.sha256('\n'.join(rows).encode()).hexdigest())


def paths(cells, start):
    routes = {start: []}
    queue = deque([start])
    while queue:
        x, y = queue.popleft()
        for a, (dx, dy) in enumerate(DELTAS):
            q = (x+dx, y+dy)
            if q in cells and q not in routes:
                routes[q] = routes[(x, y)] + [a]
                queue.append(q)
    return routes


class World:
    def __init__(self, spec, horizon):
        self.rows = spec['rows']
        self.pos = tuple(spec['start'])
        self.goal = tuple(spec['goal'])
        self.horizon, self.t, self.ended = horizon, 0, False

    def observation(self):
        x, y = self.pos
        # Local categorical sensor, no ray casting, perfect absolute localization.
        patch = [[1 if self.rows[y+dy][x+dx] == '#' else
                  (2 if (x+dx, y+dy) == self.goal else 0)
                  for dx in (-1, 0, 1)] for dy in (-1, 0, 1)]
        return dict(position=list(self.pos), goal=list(self.goal), patch=patch)

    def step(self, action):
        if self.ended:
            raise RuntimeError('reset required after episode end')
        if type(action) is not int or not 0 <= action < 4:
            raise ValueError('action must be integer 0..3')
        dx, dy = DELTAS[action]
        q = (self.pos[0]+dx, self.pos[1]+dy)
        collision = self.rows[q[1]][q[0]] == '#'
        if not collision:
            self.pos = q
        self.t += 1
        terminated = self.pos == self.goal
        truncated = self.t >= self.horizon  # May coexist with terminated.
        reward = float(terminated) - 0.01 - 0.05 * collision
        self.ended = terminated or truncated
        return self.observation(), reward, terminated, truncated, dict(collision=collision)


class Policy:
    def __init__(self, kind, seed):
        self.kind, self.rng = kind, random.Random(seed)
        self.known, self.visited = {}, set()

    def act(self, obs):
        # Intentionally receives only the public observation, never world/info/spec.
        if self.kind == 'random':
            return self.rng.randrange(4)
        if self.kind == 'memory_reset':
            self.known, self.visited = {}, set()
        pos, goal = tuple(obs['position']), tuple(obs['goal'])
        for iy, row in enumerate(obs['patch']):
            for ix, value in enumerate(row):
                self.known[(pos[0]+ix-1, pos[1]+iy-1)] = value
        self.visited.add(pos)
        routes = paths({p for p, v in self.known.items() if v != 1}, pos)
        if goal in routes:
            return routes[goal][0]
        candidates = [p for p in routes if p not in self.visited]
        if not candidates:
            raise RuntimeError('reachable unknown frontier exhausted without goal')
        target = min(candidates, key=lambda p: (len(routes[p]),
                     abs(p[0]-goal[0])+abs(p[1]-goal[1]), p[1], p[0]))
        return routes[target][0]


def episode(spec, kind, horizon):
    env = World(spec, horizon)
    policy = Policy(kind, 73_000 + spec['seed'])
    cells = {(x, y) for y, row in enumerate(spec['rows'])
             for x, v in enumerate(row) if v == '.'}
    shortest = len(paths(cells, env.pos)[env.goal])
    assert shortest > 0
    obs, trace, total, moved, collisions = env.observation(), [], 0., 0, 0
    while not env.ended:
        if kind == 'oracle':
            action = paths(cells, env.pos)[env.goal][0]  # Privileged diagnostic only.
        else:
            action = policy.act(obs)
        before = list(env.pos)
        after, reward, term, trunc, info = env.step(action)
        collisions += info['collision']
        moved += before != after['position']
        total += reward
        trace.append(dict(t=env.t, observation=obs, action=action,
                          next_observation=after, reward=reward, terminated=term,
                          truncated=trunc, collision=info['collision']))
        obs = after
    success = env.pos == env.goal
    return dict(seed=spec['seed'], policy=kind, horizon=horizon,
                map_sha256=spec['map_sha256'], shortest=shortest,
                success=success, steps=env.t, moves=moved, collisions=collisions,
                collision_rate=collisions/env.t, return_=total,
                spl=success*shortest/max(shortest, moved),
                action_efficiency=success*shortest/max(shortest, env.t),
                trace=trace)


def aggregate(runs):
    return {field: statistics.mean(float(r[field]) for r in runs)
            for field in ('success', 'steps', 'moves', 'collision_rate', 'return_',
                          'spl', 'action_efficiency')}


def checks():
    spec = dict(rows=['#####', '#...#', '#####'], start=[1, 1], goal=[2, 1])
    env = World(spec, 1)
    _, r, term, trunc, _ = env.step(1)
    assert term and trunc and r == .99
    try:
        env.step(1)
        raise AssertionError('post-terminal step accepted')
    except RuntimeError:
        pass
    env = World(spec, 1)
    _, r, term, trunc, info = env.step(0)
    assert not term and trunc and info['collision'] and abs(r + .06) < 1e-12
    env = World(spec, 8)
    for bad in (-1, 4, 'east', True):
        try:
            env.step(bad)
            raise AssertionError('invalid action accepted')
        except ValueError:
            pass
    env = World(spec, 3)
    for a in (0, 0, 1):
        _, _, term, trunc, _ = env.step(a)
    assert term and trunc and env.pos == (2, 1) and env.t == 3
    # One-cell path gives SPL=1; all three actions give efficiency=1/3.
    # A hypothetical value target, not a learned Q function.
    correct_target = -.06 + .99 * .5
    assert abs(correct_target - .435) < 1e-12
    return dict(terminal_at_limit='both flags true', wall='stationary, penalty, timeout',
                invalid_actions=4, post_terminal='rejected', collision_metric_fixture='two blocked actions then one successful move',
                truncated_bootstrap_target=correct_target, wrong_done_target=-.06)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--out', type=Path, default=ROOT/'results')
    args = parser.parse_args()
    seen, rejected = set(), []
    def unique_maps(first_seed, count):
        selected, seed = [], first_seed
        while len(selected) < count:
            spec = maze(seed)
            if spec['map_sha256'] in seen:
                rejected.append(seed)
            else:
                selected.append(spec)
                seen.add(spec['map_sha256'])
            seed += 1
        return selected
    dev = unique_maps(73, 5)
    test = unique_maps(100, 30)
    assert len({s['map_sha256'] for s in dev+test}) == 35
    dump(args.out/'maps.json', dict(dev=dev, test=test, duplicate_seeds_rejected=rejected))
    dump(args.out/'checks.json', checks())
    all_runs, summary = [], {}
    for horizon in (16, 32, 64):
        summary[str(horizon)] = {}
        for kind in POLICIES:
            runs = [episode(s, kind, horizon) for s in test]
            all_runs.extend(runs)
            summary[str(horizon)][kind] = aggregate(runs)
    with (args.out/'episodes.jsonl').open('w') as f:
        for run in all_runs:
            f.write(json.dumps(run) + '\n')
    dump(args.out/'summary.json', summary)
    print('Python', sys.version.split()[0], '| standard library CPU | map [9,9], patch [3,3], action scalar [0..3]')
    print('35 unique maps: 5 reserved development, 30 test; no training or tuning')
    print('360 episodes saved; random policy seed = 73000 + map seed')
    print(json.dumps(summary, indent=2))
    print('Boundary checks PASSED:', json.dumps(checks()))


if __name__ == '__main__':
    main()