#!/usr/bin/env python3
"""Reconstruct decisions from privately held dataset snapshots and a case manifest."""
import argparse,csv,json
from pathlib import Path


def main():
    p=argparse.ArgumentParser();p.add_argument('--snapshots',type=Path,required=True);p.add_argument('--manifest',type=Path,required=True)
    p.add_argument('--sources',type=Path,required=True);p.add_argument('--output',type=Path,required=True);a=p.parse_args()
    manifest=json.loads(a.manifest.read_text());sources={r['domain']:r for r in json.loads(a.sources.read_text())};cache={}
    with a.output.open('x') as out:
        for item in manifest['case_identifiers']['public']:
            domain=item['domain'];index=int(item['id'].rsplit('-',1)[1])
            if domain=='banking20':
                if domain not in cache:
                    with (a.snapshots/'banking-test.csv').open() as f:cache[domain]=list(csv.DictReader(f))
                r=cache[domain][index];labels=sources[domain]['classes']
                state=r['text'];question='Which banking support intent best matches the customer message?'
                options=[[x,x.replace('_',' ')] for x in labels];gold=r['category']
            else:
                if domain not in cache:
                    data=json.loads((a.snapshots/(domain+'-rows.json')).read_text())
                    cache[domain]=({x['row_idx']:x['row'] for x in data['rows']},next((f['type']['names'] for f in data['features'] if f['name']=='label'),None))
                rows,names=cache[domain];r=rows[index]
                if domain in ['news','emotion']:
                    state=r['text'];question='Which topic best describes this news item?' if domain=='news' else 'Which emotion is predominantly expressed by the author?'
                    options=[[n,n] for n in names];gold=names[r['label']]
                elif domain=='reading':
                    state=r['passage'];question='Based on the passage, answer yes or no: '+r['question'];options=[['yes','Yes'],['no','No']];gold='yes' if r['answer'] else 'no'
                elif domain=='entailment':
                    state='Premise: '+r['premise']+'\nHypothesis: '+r['hypothesis'];question='Given only the premise, is the hypothesis entailed, contradicted, or neither (neutral)?'
                    options=[['entailment','The premise entails the hypothesis'],['neutral','Neither entailed nor contradicted by the premise'],['contradiction','The premise contradicts the hypothesis']];gold=names[r['label']]
                elif domain=='science':
                    state='Choose the best answer to this science question.';question=r['question'];options=list(map(list,zip(r['choices']['label'],r['choices']['text'])));gold=r['answerKey']
                else:raise ValueError('Unknown domain '+domain)
            out.write(json.dumps(dict(id=item['id'],request=dict(state=state,question=question,options=options),expected=gold))+'\n')


if __name__=='__main__':main()