"""Original LearnAI CPU Q-learning reference. Python 3.10+, standard library only.
This is a small public grid experiment, not a standard robotics simulator or P09 completion.
"""
import argparse
import copy
import json
import math
from pathlib import Path

SIZES = {'w1': 72, 'b1': 8, 'w2': 32, 'b2': 4}
SEEDS = (7, 19, 41)
DEFAULT = {'replay': True, 'targetEvery': 20, 'rate': .01, 'epsilon': .3}

def integer(v, lo, hi):
    return type(v) is int and lo <= v <= hi

def config_check(v):
    if not isinstance(v, dict) or type(v.get('replay')) is not bool or type(v.get('targetEvery')) is not int or v.get('targetEvery') not in (1, 20) or type(v.get('rate')) not in (float, int) or v.get('rate') not in (.01, .05) or type(v.get('epsilon')) not in (float, int) or v.get('epsilon') not in (.1, .3):
        raise ValueError('Invalid bounded configuration')
    return {k: v[k] for k in DEFAULT}

def transition(state, action):
    if not integer(state, 0, 8) or state in (5, 8) or not integer(action, 0, 3):
        raise ValueError('A command needs a nonterminal cell and action 0..3')
    row, col = divmod(state, 3)
    dy, dx = ((-1, 0), (0, 1), (1, 0), (0, -1))[action]
    r, c = row + dy, col + dx
    nxt = r * 3 + c if 0 <= r < 3 and 0 <= c < 3 else state
    return dict(state=state, action=action, next=nxt, reward=10 if nxt == 8 else -10 if nxt == 5 else -1, terminal=nxt in (5, 8))

def network(seed):
    if seed not in SEEDS or type(seed) is not int:
        raise ValueError('Choose a declared seed: 7, 19 or 41')
    return dict(w1=[math.sin((i+1)*12.9898+seed)*.1 for i in range(72)], b1=[0.]*8,
                w2=[math.sin((i+1)*7.13+seed)*.1 for i in range(32)], b2=[0.]*4)

def initial(seed=7):
    online = network(seed)
    return dict(online=online, target=copy.deepcopy(online), buffer=[], updates=0, random=seed, position=0,
                episodeSteps=0, episodeReturn=0, episodes=0, successes=0, truncations=0, lastLoss=None)

def forward(n, state):
    if not integer(state, 0, 8):
        raise ValueError('Invalid cell')
    h = [math.tanh(n['b1'][j]+n['w1'][j*9+state]) for j in range(8)]
    q = [n['b2'][a]+sum(h[j]*n['w2'][a*8+j] for j in range(8)) for a in range(4)]
    return h, q

def stopped_target(receipt, next_values):
    """Task termination suppresses bootstrap; collection timeout is not termination."""
    return receipt['reward'] + (0 if receipt['terminal'] else .9*max(next_values))

def gradient(online, target, batch):
    if not 1 <= len(batch) <= 8:
        raise ValueError('One to eight receipts required')
    g = {k: [0.]*size for k, size in SIZES.items()}
    loss = 0.
    for receipt in batch:
        if not isinstance(receipt, dict) or type(receipt.get('next')) is not int or type(receipt.get('reward')) is not int or type(receipt.get('terminal')) is not bool or receipt != transition(receipt.get('state'), receipt.get('action')):
            raise ValueError('Receipt violates the training simulator')
        state, action = receipt['state'], receipt['action']
        h, q = forward(online, state)
        y = stopped_target(receipt, forward(target, receipt['next'])[1])
        error = q[action] - y
        loss += error*error/2/len(batch)
        d = error/len(batch)
        g['b2'][action] += d
        for j in range(8):
            g['w2'][action*8+j] += d*h[j]
            z = d*online['w2'][action*8+j]*(1-h[j]*h[j])
            g['b1'][j] += z
            g['w1'][j*9+state] += z
    return g, loss

def validate(c):
    if not isinstance(c, dict):
        raise ValueError('Invalid checkpoint')
    for name in ('online', 'target'):
        n = c.get(name)
        if not isinstance(n, dict):
            raise ValueError('Invalid network')
        for key, size in SIZES.items():
            values = n.get(key)
            if not isinstance(values, list) or len(values) != size or any(type(v) not in (int, float) or not math.isfinite(v) or abs(v) > 1000 for v in values):
                raise ValueError('Invalid network parameters')
    for key, upper in [('updates', 2000), ('random', 2**32-1), ('position', 8), ('episodeSteps', 23), ('episodes', 2000), ('successes', 2000), ('truncations', 2000)]:
        if not integer(c.get(key), 0, upper):
            raise ValueError('Invalid counter')
    if c['position'] in (5, 8) or c['episodes'] > c['updates'] or c['successes']+c['truncations'] > c['episodes'] or c['episodeSteps'] > c['updates'] or type(c.get('episodeReturn')) is not int or c['episodeReturn'] != -c['episodeSteps'] or (c['episodeSteps'] == 0 and c['position'] != 0):
        raise ValueError('Invalid episode state')
    loss = c.get('lastLoss')
    if (c['updates'] == 0) != (loss is None) or (loss is not None and (type(loss) not in (float, int) or not math.isfinite(loss) or loss < 0)):
        raise ValueError('Invalid loss')
    if not isinstance(c.get('buffer'), list) or len(c['buffer']) != min(64, c['updates']):
        raise ValueError('Invalid replay length')
    for t in c['buffer']:
        if not isinstance(t, dict) or type(t.get('terminal')) is not bool or type(t.get('next')) is not int or type(t.get('reward')) is not int or t != transition(t.get('state'), t.get('action')):
            raise ValueError('Invalid replay receipt')
    return copy.deepcopy(c)

def train(prior, config, steps=1):
    config = config_check(config)
    c = validate(prior)
    if not integer(steps, 1, 25) or c['updates']+steps > 2000:
        raise ValueError('Calls cap at 25 updates; total cap 2000')
    def random():
        c['random'] = (c['random']*1664525+1013904223) % 2**32
        return c['random']/2**32
    for _ in range(steps):
        q = forward(c['online'], c['position'])[1]
        action = int(random()*4) if random() < config['epsilon'] else max(range(4), key=lambda a: q[a])
        receipt = transition(c['position'], action)
        c['buffer'] = (c['buffer']+[receipt])[-64:]
        batch = []
        for _ in range(8):
            u = random()  # Consume the same draws in newest-only mode.
            batch.append(c['buffer'][int(u*len(c['buffer']))] if config['replay'] else receipt)
        g, loss = gradient(c['online'], c['target'], batch)
        candidate = {k: [w-config['rate']*d for w, d in zip(c['online'][k], g[k])] for k in SIZES}
        if any(not math.isfinite(w) or abs(w) > 1000 for arr in candidate.values() for w in arr):
            raise ValueError('Rejected proposal; caller retains preceding checkpoint')
        c['online'] = candidate
        c['updates'] += 1
        c['lastLoss'] = loss
        if c['updates'] % config['targetEvery'] == 0:
            c['target'] = copy.deepcopy(candidate)
        c['position'] = receipt['next']
        c['episodeSteps'] += 1
        c['episodeReturn'] += receipt['reward']
        if receipt['terminal'] or c['episodeSteps'] == 24:
            c['episodes'] += 1
            c['successes'] += int(receipt['next'] == 8)
            c['truncations'] += int(not receipt['terminal'])
            c['position'] = c['episodeSteps'] = c['episodeReturn'] = 0
    return validate(c)

def read_artifact(path):
    path = Path(path)
    if path.stat().st_size > 500000:
        raise ValueError('Artifact exceeds 500,000 bytes')
    def reject_constant(v):
        raise ValueError('Nonfinite JSON constant: '+v)
    v = json.loads(path.read_text(), parse_constant=reject_constant)
    if not isinstance(v, dict) or v.get('format') != 'dqn-practice-v1' or not isinstance(v.get('current'), dict):
        raise ValueError('Use a Q-learning lesson/studio artifact')
    current = v['current']
    if type(current.get('seed')) is not int or current['seed'] not in SEEDS:
        raise ValueError('Invalid seed')
    return dict(seed=current['seed'], config=config_check(current.get('config')), checkpoint=validate(current.get('checkpoint')))

def artifact(current):
    return dict(format='dqn-practice-v1', current=current, humanReview='pending',
                limits='Public toy-grid practice; counters and authorship are not authenticated. Import numeric models into the browser project, freeze three matched seeds and explicitly open evaluation. Python results do not award proficiency.')

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='command', required=True)
    fresh = sub.add_parser('train')
    fresh.add_argument('output', type=Path)
    fresh.add_argument('--seed', type=int, choices=SEEDS, default=7)
    fresh.add_argument('--replay', choices=['on', 'off'], default='on')
    fresh.add_argument('--target-every', type=int, choices=[1, 20], default=20)
    fresh.add_argument('--rate', type=float, choices=[.01, .05], default=.01)
    fresh.add_argument('--epsilon', type=float, choices=[.1, .3], default=.3)
    resume = sub.add_parser('resume')
    resume.add_argument('input', type=Path)
    resume.add_argument('output', type=Path)
    for command in [fresh, resume]:
        command.add_argument('--steps', type=int, default=200)
    args = parser.parse_args()
    if not integer(args.steps, 1, 2000):
        parser.error('Choose 1..2000 additional updates')
    if args.output.exists():
        parser.error('Output already exists; choose a new file')
    current = read_artifact(args.input) if args.command == 'resume' else dict(seed=args.seed, config=dict(replay=args.replay=='on', targetEvery=args.target_every, rate=args.rate, epsilon=args.epsilon), checkpoint=initial(args.seed))
    if current['checkpoint']['updates']+args.steps > 2000:
        parser.error('Total update cap is 2000')
    remaining = args.steps
    while remaining:
        count = min(25, remaining)
        current['checkpoint'] = train(current['checkpoint'], current['config'], count)
        remaining -= count
    encoded = json.dumps(artifact(current), indent=2, allow_nan=False)+'\n'
    with args.output.open('x') as output:
        output.write(encoded)
    print(f"Saved {args.output}: {current['checkpoint']['updates']} updates; no reserved evaluation opened")

if __name__ == '__main__':
    main()