"""Public split and canonical-boundary checks; does not decode or pair DICOMs."""
from __future__ import annotations

import argparse
from collections import defaultdict
import hashlib
import json
from pathlib import Path, PurePosixPath
import re
from statistics import fmean

DEVICES = {'heidelberg_spectralis', 'topcon_maestro2', 'topcon_triton', 'zeiss_cirrus'}
SEED = 20261105


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def relative_path(value):
    if not isinstance(value, str) or not value or '\\' in value or '\x00' in value or re.match(r'^[A-Za-z]:', value):
        raise ValueError('Dataset paths must be relative POSIX paths')
    path = PurePosixPath(value)
    if path.is_absolute() or '..' in path.parts or value == '.':
        raise ValueError('Dataset paths must stay beneath the supplied data root')
    return value


def load_pairs(path):
    rows, identities, inputs = [], set(), set()
    for line in Path(path).read_text().splitlines():
        if not line.strip():
            continue
        row = json.loads(line)
        if not isinstance(row, dict):
            raise ValueError('Manifest rows must be objects')
        name = row.get('case_id')
        if not isinstance(name, str) or not re.fullmatch(r'[A-Za-z0-9_-]{1,64}', name) or name in identities:
            raise ValueError('Case IDs must be unique safe filename stems')
        if (not isinstance(row.get('participant_id'), str) or not row['participant_id'].strip()
                or row['participant_id'] != row['participant_id'].strip()):
            raise ValueError('Each case needs a nonempty participant identity')
        if row.get('device') not in DEVICES:
            raise ValueError('Device family must match the declared AI-READI cohort')
        for field in ('oct_paths', 'segmentation_paths'):
            values = row.get(field)
            if not isinstance(values, list) or not values:
                raise ValueError('Each case needs structural input and segmentation paths')
            checked = [relative_path(item) for item in values]
            if len(set(checked)) != len(checked):
                raise ValueError('Duplicate paths within a case')
        if inputs.intersection(row['oct_paths']):
            raise ValueError('A structural input cannot be duplicated across cases')
        inputs.update(row['oct_paths'])
        identities.add(name)
        rows.append(row)
    if not rows:
        raise ValueError('The manifest is empty')
    return rows


def split_for(participant):
    # Integer buckets avoid floating-point or Python-hash dependence.
    key = f'aireadi-oct-v1:{SEED}:{participant}'.encode('utf-8')
    bucket = int.from_bytes(hashlib.sha256(key).digest(), 'big') % 100
    return 'train' if bucket < 80 else 'validation' if bucket < 90 else 'test'


def write_splits(pairs, output):
    rows = load_pairs(pairs)
    groups = {name: [] for name in ('train', 'validation', 'test')}
    for row in sorted(rows, key=lambda item: item['case_id']):
        groups[split_for(row['participant_id'])].append(row)
    if any(not group for group in groups.values()):
        raise ValueError('A split is empty; report this rather than rerolling the seed')
    output = Path(output)
    output.mkdir(parents=True, exist_ok=False, mode=0o700)
    report = {'seed': SEED, 'algorithm': 'sha256-mod100-v1', 'ratios': [80, 10, 10],
              'source_manifest_sha256': digest(pairs), 'splits': {}}
    for name, group in groups.items():
        destination = output / (name + '.jsonl')
        destination.write_text(''.join(json.dumps(row, sort_keys=True) + '\n' for row in group))
        public = [{key: row[key] for key in ('case_id', 'device', 'oct_paths')} for row in group]
        (output / (name + '-inputs.jsonl')).write_text(''.join(json.dumps(row) + '\n' for row in public))
        report['splits'][name] = {'cases': len(group), 'participants': len({row['participant_id'] for row in group}),
            'devices': sorted({row['device'] for row in group}), 'manifest_sha256': digest(destination)}
    (output / 'split-summary.json').write_text(json.dumps(report, indent=2) + '\n')
    return report


def read_arrays(path, reference=False):
    import numpy as np
    with np.load(path, allow_pickle=False) as data:
        required = {'positions', 'boundary_keys', 'valid', 'height'} if reference else {'positions', 'boundary_keys'}
        allowed = required | ({'axial_spacing_um'} if reference else set())
        if not required <= set(data.files) or set(data.files) - allowed:
            raise ValueError('Unexpected or missing canonical array fields')
        positions = data['positions']
        keys = data['boundary_keys']
        if positions.ndim != 3 or any(n == 0 for n in positions.shape) or positions.dtype.kind not in 'fiu':
            raise ValueError('positions must be a nonempty real [boundary, B-scan, A-scan] array')
        if keys.ndim != 1 or keys.dtype.kind not in 'US' or len(keys) != positions.shape[0]:
            raise ValueError('boundary_keys must be a one-dimensional Unicode/byte-string array')
        names = [x.decode('utf-8') if isinstance(x, bytes) else str(x) for x in keys.tolist()]
        if any(not x or x != x.strip() for x in names) or len(set(names)) != len(names):
            raise ValueError('Boundary keys must be nonempty and unique')
        result = {'positions': positions.astype('float64'), 'keys': names}
        if reference:
            valid, height = data['valid'], data['height']
            if valid.dtype.kind != 'b' or valid.shape != positions.shape:
                raise ValueError('Reference valid must be a matching Boolean mask')
            if height.shape != () or height.dtype.kind not in 'iu' or int(height) < 2:
                raise ValueError('Reference height must be the original axial size')
            if not valid.any() or not np.isfinite(positions[valid]).all():
                raise ValueError('Reference needs finite valid annotation support')
            if ((positions[valid] < 0) | (positions[valid] > int(height) - 1)).any():
                raise ValueError('Reference coordinates are outside original geometry')
            result.update(valid=valid.copy(), height=int(height), spacing=None)
            if 'axial_spacing_um' in data:
                spacing = data['axial_spacing_um']
                if spacing.shape != () or spacing.dtype.kind not in 'fiu' or not np.isfinite(spacing) or float(spacing) <= 0:
                    raise ValueError('Axial spacing must be a positive finite scalar in micrometres')
                result['spacing'] = float(spacing)
        return result


def score(manifest, references, predictions, expected_manifest_sha256=None):
    import numpy as np
    fingerprint = digest(manifest)
    if expected_manifest_sha256 is not None and fingerprint != expected_manifest_sha256:
        raise ValueError('The frozen manifest changed')
    rows = load_pairs(manifest)
    expected = {row['case_id'] for row in rows}
    if {path.stem for path in Path(predictions).glob('*.npz')} != expected:
        raise ValueError('Prediction cases do not exactly cover the manifest')
    grouped = defaultdict(lambda: defaultdict(list))
    boundaries = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    boundary_counts = defaultdict(lambda: defaultdict(int))
    physical = defaultdict(lambda: defaultdict(list))
    spacing_cases = defaultdict(int)
    for row in rows:
        name = row['case_id'] + '.npz'
        ref = read_arrays(Path(references) / name, reference=True)
        pred = read_arrays(Path(predictions) / name)
        if set(ref['keys']) != set(pred['keys']) or ref['positions'].shape != pred['positions'].shape:
            raise ValueError('Prediction boundary identities or geometry do not match')
        values = pred['positions'][[pred['keys'].index(key) for key in ref['keys']]]
        if not np.isfinite(values).all() or ((values < 0) | (values > ref['height'] - 1)).any():
            raise ValueError('All predicted coordinates must be finite and in original bounds')
        device, participant = row['device'], row['participant_id']
        errors = []
        for index, key in enumerate(ref['keys']):
            mask = ref['valid'][index]
            boundary_counts[device][key] += int(mask.sum())
            if mask.any():
                error = float(np.abs(values[index][mask] - ref['positions'][index][mask]).mean())
                errors.append(error)
                boundaries[device][key][participant].append(error)
        case_error = fmean(errors)
        grouped[device][participant].append(case_error)
        if ref['spacing'] is not None:
            physical[device][participant].append(case_error * ref['spacing'])
            spacing_cases[device] += 1
    average = lambda participants: fmean(fmean(cases) for cases in participants.values())
    devices = {}
    for device, participants in sorted(grouped.items()):
        count = sum(map(len, participants.values()))
        devices[device] = {'cases': count, 'participants': len(participants), 'mae_pixels': average(participants),
            'spacing_cases': spacing_cases[device],
            'mae_um': average(physical[device]) if spacing_cases[device] == count else None,
            'boundaries': {key: {'mae_pixels': average(groups) if groups else None, 'participants': len(groups),
                                'valid_points': boundary_counts[device][key]}
                           for key, groups in ((key, boundaries[device][key]) for key in sorted(boundary_counts[device]))}}
    return {'measurement': 'informal_canonical_boundary_agreement', 'manifest_sha256': fingerprint,
        'cases': len(rows), 'participants': len({row['participant_id'] for row in rows}),
        'mae_pixels': fmean(item['mae_pixels'] for item in devices.values()),
        'mae_um': fmean(item['mae_um'] for item in devices.values()) if all(item['mae_um'] is not None for item in devices.values()) else None,
        'aggregation': 'valid points -> boundaries -> volume -> participant/device -> device -> macro device',
        'devices': devices}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest='command', required=True)
    split = commands.add_parser('split')
    split.add_argument('--pairs', required=True, type=Path)
    split.add_argument('--output-dir', required=True, type=Path)
    measure = commands.add_parser('score')
    measure.add_argument('--manifest', required=True, type=Path)
    measure.add_argument('--reference-dir', required=True, type=Path)
    measure.add_argument('--predictions-dir', required=True, type=Path)
    measure.add_argument('--expected-manifest-sha256')
    measure.add_argument('--output', required=True, type=Path)
    args = parser.parse_args()
    if args.command == 'split':
        result = write_splits(args.pairs, args.output_dir)
    else:
        if args.output.exists():
            raise FileExistsError('Choose a fresh metric report path')
        result = score(args.manifest, args.reference_dir, args.predictions_dir, args.expected_manifest_sha256)
        with args.output.open('x') as stream:
            json.dump(result, stream, indent=2, allow_nan=False)
            stream.write('\n')
    print(json.dumps(result, allow_nan=False))


if __name__ == '__main__':
    main()