# -*- coding: utf-8 -*-

from __future__ import print_function
import argparse
import datetime
import fileinput
import math
import os
import os.path
import shelve
import sys

import numpy
import pandas

import pig_schema

from restrictions import MeanRankInRangeAtDate

from midas.scripts import MDCommand

SCHEMA = pig_schema.pig_schema_to_py_struct(
    ','.join(['(site: chararray',
              'ranking: bag{(tstamp: chararray, rank: int)}',
              'company: chararray',
              'code: chararray',
              'tstamp: chararray)'])
    )
PARSER = pig_schema.make_parser(SCHEMA)

ARGPARSER = argparse.ArgumentParser()
ARGPARSER.add_argument('files', metavar='FILE', nargs='+', 
                       help='The `sites_w_company` files to extract restrictions from.')

DATE_INTERVAL = pandas.DateOffset(days=3)
OFFSET = pandas.DateOffset(days=90)

def get_median_rank_at_funding_date(field):
    fund_date = pandas.Timestamp(field.tstamp) - OFFSET
    s = pandas.Series(dict((pandas.Timestamp(d), r)
                           for d, r in field.ranking))
    r = s[fund_date-DATE_INTERVAL:fund_date+DATE_INTERVAL].median()
    return fund_date, r

def main(iterator):
    for i in iterator:
        field = PARSER(i)
        fund_date, mean_rank = get_median_rank_at_funding_date(field)
        if numpy.isnan(mean_rank):
            continue
        mean_rank = int(mean_rank)
        border = int(2 ** math.log(mean_rank))
        restr = MeanRankInRangeAtDate(fund_date, 
                                      mean_rank - border,
                                      mean_rank + border)
        SHELF[field.company] = restr
        

if __name__ == '__main__':
    args = ARGPARSER.parse_args()
    SHELF = shelve.open('restrictions_shelve')
    main(fileinput.input(args.files))
    SHELF.sync()
    SHELF.close()


@@ -67,83 +67,25 @@
 ort 
-argparse%0Aimport datetime%0Aimport fileinput%0Aimport math%0Aimport os%0Aimport os.p
+datetime%0Aimport m
 ath%0A
@@ -101,19 +101,8 @@
 elve
-%0Aimport sys
 %0A%0Aim
@@ -131,32 +131,113 @@
 as%0A%0A
-import pig_schema%0A%0Afrom 
+from midas.pig_schema import pig_schema_to_py_struct%0Afrom midas.pig_schema import make_parser%0Afrom midas.
 rest
@@ -274,17 +274,16 @@
 eAtDate%0A
-%0A
 from mid
@@ -314,29 +314,77 @@
 and%0A
-%0ASCHEMA = pig_schema.
+from midas.scripts import StoreSingleFileOrDirectoryAction%0A%0ASCHEMA = 
 pig_
@@ -625,19 +625,8 @@
 R = 
-pig_schema.
 make
@@ -646,196 +646,8 @@
 A)%0A%0A
-ARGPARSER = argparse.ArgumentParser()%0AARGPARSER.add_argument('files', metavar='FILE', nargs='+', %0A                       help='The %60sites_w_company%60 files to extract restrictions from.')%0A%0A
 DATE
@@ -1025,59 +1025,535 @@
  r%0A%0A
-def main(iterator):%0A    for i in iterator:%0A
+class MakeSampleRestrictionShelve(MDCommand):%0A%0A    def add_argument(self):%0A        self.parser.add_argument(%0A            '--shelf', metavar='FILE',%0A            help='The file-name of the shelve.'%0A            )%0A        self.parser.add_argument(%0A            'sites_w_company', action=StoreSingleFileOrDirectoryAction,%0A            help='The directory or file where sites_w_company data resides.'%0A            )%0A%0A    def run(self):%0A        shelf = shelve.open(self.args.shelf)%0A        for i in self.args.sites_w_company:%0A   
         
 fiel
@@ -1548,16 +1548,17 @@
         
+ 
 field = 
@@ -1567,16 +1567,20 @@
 RSER(i)%0A
+    
         
@@ -1649,16 +1649,20 @@
         
+   
+ 
 if numpy
@@ -1692,16 +1692,20 @@
         
+    
 continue
@@ -1705,32 +1705,36 @@
 ontinue%0A        
+   
+ 
 mean_rank = int(
@@ -1744,16 +1744,20 @@
 n_rank)%0A
+    
         
@@ -1795,16 +1795,20 @@
 _rank))%0A
+    
         
@@ -1879,32 +1879,36 @@
                 
+    
 mean_rank - bord
@@ -1911,16 +1911,20 @@
 border,%0A
+    
         
@@ -1981,21 +1981,25 @@
         
-SHELF
+    shelf
 %5Bfield.c
@@ -2026,182 +2026,34 @@
     
-%0A%0Aif __name__ == '__main__':%0A    args = ARGPARSER.parse_args()%0A    SHELF = shelve.open('restrictions_shelve')%0A    main(fileinput.input(args.files))%0A    SHELF.sync()%0A    SHELF
+shelf.sync()%0A        shelf
 .clo
@@ -2057,9 +2057,8 @@
 close()%0A
-%0A

rewrite to use current library