"""P06 supplied 3x3 RGB convolutional segmenter. Python 3.11+, CPU, no dependencies.
Two shared tanh filters followed by a shared logistic head: 59 learned parameters.
Public development only; independently authored extensions/review remain pending.
"""
import argparse
import importlib.util
import json
import math
from pathlib import Path

spec = importlib.util.spec_from_file_location('pixel_baseline', Path(__file__).with_name('perception-baseline.py'))
b = importlib.util.module_from_spec(spec)
spec.loader.exec_module(b)
VERSION = 'p06-convolution-v1'


def initial():
    return [0.08 * math.sin(i + 1) for i in range(54)] + [0., 0., .2, -.2, 0.]


def validate(weights):
    if not isinstance(weights, list) or len(weights) != 59 or any(
            type(w) not in (int, float) or not math.isfinite(w) or abs(w) > 50 for w in weights):
        raise ValueError('Expected 59 finite weights bounded by 50.')
    return weights


def patch(record, index):
    """Row-major 3x3 neighborhood, RGB innermost; zero padding in normalized space.
    Padding zero represents mid-gray, not black. Label-3 pixels remain valid RGB
    context, but their own outputs are excluded from supervised loss and metrics.
    """
    width, height = record['width'], record['height']
    x, y = index % width, index // width
    result = []
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            xx, yy = x + dx, y + dy
            if 0 <= xx < width and 0 <= yy < height:
                start = 3 * (yy * width + xx)
                result.extend(v / 127.5 - 1 for v in record['rgb'][start:start+3])
            else:
                result.extend([0., 0., 0.])
    return result


def forward(weights, x):
    h = [math.tanh(sum(weights[27*c+j]*v for j, v in enumerate(x)) + weights[54+c]) for c in range(2)]
    z = weights[56]*h[0] + weights[57]*h[1] + weights[58]
    return h, z, b.sigmoid(z)


def rows_for(records, split):
    return [(patch(r, i), int(label == 1)) for r in records if r['split'] == split
            for i, label in enumerate(r['trimap']) if label != 3]


def objective(weights, rows):
    validate(weights)
    if not rows:
        raise ValueError('No classified pixels.')
    gradient, loss = [0.] * 59, 0.
    for x, y in rows:
        h, z, p = forward(weights, x)
        loss += max(z, 0) - z*y + math.log1p(math.exp(-abs(z)))
        dz = (p-y) / len(rows)
        for c in range(2):
            dh = dz * weights[56+c] * (1-h[c]*h[c])
            for j in range(27):
                gradient[27*c+j] += dh * x[j]
            gradient[54+c] += dh
            gradient[56+c] += dz * h[c]
        gradient[58] += dz
    return loss / len(rows), gradient


def train(records, updates=10, rate=.3):
    if type(updates) is not int or not 1 <= updates <= 100:
        raise ValueError('Use 1–100 full-batch CPU updates.')
    if type(rate) not in (int, float) or not math.isfinite(rate) or not 0 < rate <= 1:
        raise ValueError('Use a finite learning rate above zero and at most one.')
    rows = rows_for(records, 'training')
    weights = initial()
    before = objective(weights, rows)[0]
    for _ in range(updates):
        _, gradient = objective(weights, rows)
        weights = validate([w-rate*g for w, g in zip(weights, gradient)])
    return weights, {'initialTrainingLoss': before, 'finalTrainingLoss': objective(weights, rows)[0]}


def replay(data, artifact):
    if artifact.get('version') != VERSION or artifact.get('sourceHash') != b.SOURCE_HASH:
        raise ValueError('Wrong convolution artifact or source hash.')
    weights = validate(artifact['weights'])
    threshold = artifact['threshold']
    if type(threshold) not in (int, float) or not math.isfinite(threshold) or not 0 < threshold < 1:
        raise ValueError('Threshold must be strictly between zero and one.')
    images = []
    for record in data['records']:
        mask = [int(forward(weights, patch(record, i))[2] >= threshold) for i in range(len(record['trimap']))]
        images.append({'id': record['id'], 'split': record['split'], 'model': b.score(record['trimap'], mask),
                       'alwaysBackground': b.score(record['trimap'], [False]*len(mask)),
                       'alwaysPet': b.score(record['trimap'], [True]*len(mask)), 'predictedMask': mask})
    return {'version': VERSION, 'sourceHash': b.SOURCE_HASH, 'weights': weights, 'threshold': threshold,
            'audit': b.audit(data['records']), 'images': images,
            'trainingLossRecomputed': objective(weights, rows_for(data['records'], 'training'))[0],
            'architecture': '3x3 RGB same-padded convolution, two tanh channels, logistic 1x1 head; 59 parameters',
            'padding': 'zero in normalized RGB coordinates (mid-gray)',
            'evidence': 'Supplied model; public repeated development; weights do not authenticate training history.',
            'freshSourceEvaluation': 'pending', 'independentReview': 'pending'}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--data', type=Path, default=Path(__file__).with_name('subset.json'))
    parser.add_argument('--updates', type=int, default=10)
    parser.add_argument('--rate', type=float, default=.3)
    parser.add_argument('--threshold', type=float, default=.5)
    parser.add_argument('--replay', type=Path)
    parser.add_argument('--out', type=Path, required=True)
    args = parser.parse_args()
    data = b.load_data(args.data)
    if args.out.exists():
        parser.error('Output already exists. Choose a new filename.')
    if args.replay:
        with args.replay.open('rb') as stream:
            raw = stream.read(500001)
        if len(raw) > 500000:
            parser.error('Artifact exceeds 500 KB.')
        result = replay(data, json.loads(raw))
    else:
        weights, trace = train(data['records'], args.updates, args.rate)
        result = replay(data, {'version': VERSION, 'sourceHash': b.SOURCE_HASH, 'weights': weights, 'threshold': args.threshold})
        result['suppliedTrainingRun'] = {**trace, 'updates': args.updates, 'rate': args.rate}
    with args.out.open('x') as stream:
        json.dump(result, stream, indent=2, allow_nan=False)
    print('Wrote', args.out, '— public development, not independent evaluation.')


if __name__ == '__main__':
    main()