# Copyright 2017 Rice University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
import argparse
import json
import random
from itertools import chain

import bayou.models.core.evidence

HELP = """Use this script to extract evidences from a raw data file with sequences generated by driver.
You can also filter programs based on number and length of sequences, and control the samples from each program."""


def extract_evidence(clargs):
    print('Loading data file...', end='')
    with open(clargs.input_file[0]) as f:
        js = json.load(f)
    print('done')
    done = 0
    programs = []
    for program in js['programs']:
        sequences = program['sequences']
        if len(sequences) > clargs.max_seqs or \
                any([len(sequence['calls']) > clargs.max_seq_length for sequence in sequences]):
            continue

        calls = set(chain.from_iterable([sequence['calls'] for sequence in sequences]))

        apicalls = list(set(chain.from_iterable([bayou.models.core.evidence.APICalls.from_call(call) for call in calls])))
        types = list(set(chain.from_iterable([bayou.models.core.evidence.Types.from_call(call) for call in calls])))
        keywords = list(set(chain.from_iterable([bayou.models.core.evidence.Keywords.from_call(call) for call in calls])))

        if clargs.num_samples == 0:
            program['apicalls'] = apicalls
            program['types'] = types
            program['keywords'] = keywords
            programs.append(program)
        else:
            for i in range(clargs.num_samples):
                sample = dict(program)
                sample['apicalls'], sample['types'], sample['keywords'] = [], [], []
                while sample['apicalls'] == [] and sample['types'] == [] and sample['keywords'] == []:
                    sample['apicalls'] = random.sample(apicalls, random.choice(range(len(apicalls) + 1)))
                    sample['types'] = random.sample(types, random.choice(range(len(types) + 1)))
                    sample['keywords'] = random.sample(keywords, random.choice(range(len(keywords) + 1)))
                programs.append(sample)

        done += 1
        print('Extracted evidence for {} programs'.format(done), end='\r')

    print('\nWriting to {}...'.format(clargs.output_file[0]), end='')
    with open(clargs.output_file[0], 'w') as f:
        json.dump({'programs': programs}, fp=f, indent=2)
    print('done')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=HELP)
    parser.add_argument('input_file', type=str, nargs=1,
                        help='input data file')
    parser.add_argument('output_file', type=str, nargs=1,
                        help='output data file')
    parser.add_argument('--max_seqs', type=int, default=9999,
                        help='maximum number of sequences in a program')
    parser.add_argument('--max_seq_length', type=int, default=9999,
                        help='maximum length of each sequence in a program')
    parser.add_argument('--num_samples', type=int, default=0,
                        help='number of samples of evidences per program')
    clargs = parser.parse_args()
    extract_evidence(clargs)

@@ -638,16 +638,28 @@
 rt json%0A
+import math%0A
 import r
@@ -2046,16 +2046,253 @@
   else:%0A
+            # put all evidences in the same bag (to avoid bias during sampling)%0A            evidences = %5B(e, 'apicalls') for e in apicalls%5D + %5B(e, 'types') for e in types%5D + %5C%0A                        %5B(e, 'keywords') for e in keywords%5D%0A%0A
         
@@ -2408,55 +2408,10 @@
 ls'%5D
-, sample%5B'types'%5D, sample%5B'keywords'%5D = %5B%5D, %5B%5D,
+ =
  %5B%5D%0A
@@ -2430,22 +2430,16 @@
     
-while 
 sample%5B'
 apic
@@ -2438,54 +2438,36 @@
 le%5B'
-apicall
+type
 s'%5D =
-=
  %5B%5D
- and sample%5B'types'%5D == %5B%5D and
+%0A               
  sam
@@ -2487,15 +2487,12 @@
 '%5D =
-=
  %5B%5D
-:%0A 
+%0A
     
@@ -2507,298 +2507,198 @@
     
-   sample%5B'apicalls'%5D = random.sample(apicalls, random.choice(range(len(apicalls) + 1)))%0A                    sample%5B'types'%5D = random.sample(types, random.choice(range(len(types) + 1)))%0A                    sample%5B'keywords'%5D = random.sample(keywords, random.choice(range(len(keywords) + 1))
+choices = random.sample(evidences, math.ceil(len(evidences) * clargs.observability / 100))%0A                for choice, evidence in choices:%0A                    sample%5Bevidence%5D.append(choice
 )%0A  
@@ -3824,24 +3824,182 @@
 r program')%0A
+    parser.add_argument('--observability', type=int, default=100,%0A                        help='percentage of observable evidence (e.g., 100, 75, 50, etc.)')%0A
     clargs =

Add observability option to evidence_extractor.py