"""Prepare a deterministic sample with blank human labels; never submit reviews."""
import argparse
from collections import defaultdict
import csv
from hashlib import sha256
import json
from pathlib import Path
from corpus_io import DEFAULT_FOLDER, load_index, observations

FIELDS = ['observation_id', 'medicine_id', 'language', 'source_sha256', 'source_url',
          'word_paragraph_index', 'annex', 'section_code', 'before', 'after',
          'split', 'reviewer_id', 'qualification', 'label', 'medical_relevance', 'reviewed_at', 'notes']


def select(rows, count):
    groups = defaultdict(list)
    for row in rows:
        groups[row['medicine_id']].append(row)
    meds = sorted(groups, key=lambda key: sha256(key.encode()).hexdigest())
    for values in groups.values():
        values.sort(key=lambda row: row['id'])
    result = []
    position = 0
    while len(result) < count:
        available = [groups[med][position] for med in meds if len(groups[med]) > position]
        if not available:
            break
        result.extend(available[:count - len(result)])
        position += 1
    return result


def prepare(folder, destination, count):
    if count < 1:
        raise ValueError('Use a positive sample size')
    manifest, _ = load_index(folder)
    rows = select(observations(folder), count)
    destination.parent.mkdir(parents=True, exist_ok=True)
    # Refuse to overwrite an existing file of potentially completed annotations.
    with destination.open('x', encoding='utf-8', newline='') as stream:
        writer = csv.DictWriter(stream, fieldnames=FIELDS, lineterminator='\n')
        writer.writeheader()
        for row in rows:
            value = {field: '' for field in FIELDS}
            value.update({field: row.get(field, '') for field in ('medicine_id', 'language', 'source_sha256', 'source_url', 'word_paragraph_index', 'annex', 'section_code')})
            value.update(observation_id=row['id'], before=row['before']['text'], after=row['after']['text'],
                         split='evaluation' if int(sha256(row['medicine_id'].encode()).hexdigest(), 16) % 5 == 0 else 'development')
            value = {key: "'" + text if isinstance(text, str) and text.lstrip().startswith(('=', '+', '-', '@')) else text for key, text in value.items()}
            writer.writerow(value)
    return {'release_id': manifest['release_id'], 'sample_rows': len(rows),
            'output': str(destination), 'labels': 'blank', 'representative_sample': False}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--folder', type=Path, default=DEFAULT_FOLDER)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--count', type=int, default=12)
    args = parser.parse_args()
    print(json.dumps(prepare(args.folder, args.output, args.count), ensure_ascii=False, indent=2))