import os
import yaml
import ujson as json
import rpy2
from rpy2 import robjects
from rpy2.robjects.packages import importr
from collections import namedtuple

from survey_stats.log import logger
from survey_stats.survey import AnnotatedSurvey
from survey_stats.feathers import has_feather, load_feather

rfther = importr('feather', on_conflict='warn')
rutils = importr('utils')
rbase = importr('base')

cache_dir = os.path.join(os.getcwd(), 'cache')


class SurveyDataset(namedtuple('Dataset', ['config', 'surveys'])):
    __slots__ = ()

    @classmethod
    def load_dataset(cls, yml_f):
        config = None
        with open(yml_f, 'r') as fh:
            config = yaml.load(fh)['surveys']
        if not config:
            return cls(config=None, surveys=None)
        svys = {}
        for k, v in config.items():
            svys[k] = cls.fetch_or_load_dataset(k, v['spss'], v['data'])

        return cls(config=config, surveys=svys)

    def fetch_or_load_dataset(id, spss_f, data_f):
        f = os.path.join(cache_dir, '%s.feather' % id)
        ret = None
        rdf = None
        try:
            rdf = rfther.read_feather(f)
            ret = AnnotatedSurvey.from_rdf(spss_f, rdf)
            logger.info('loaded survey data from feather cache: %s' % f)
        except:
            logger.info('could not find feather cache, loading raw data')
            ret = AnnotatedSurvey.load_cdc_survey(spss_f, data_f)
            logger.info('saving data to feather cache: %s' % f)
            rfther.write_feather(ret.rdf, f)

        rbase.gc()
        return ret


class YRBSSDataset(SurveyDataset):
    __slots__ = ()

    def fetch_survey(self, combined=True, national=True, year=None):
        pred = (lambda v: v['is_combined'] == combined and
                (v['year'] == year if year else True))
        if not self.config:
            return None
        return next(( self.surveys[k] for k, v in self.config.items() if pred(v)), None)

    def fetch_config(self, national=True, year=None):
        combined = not year # if year is present, get ind year, else combined
        pred = (lambda v: v['is_combined'] == combined and
                (v['year'] == year if year else True))
        if not self.config:
            return None
        return next(((k, v) for k, v in self.config.items() if pred(v)), None)


    @property
    def survey_years(self):
        return set(sorted([int(v['year']) for v in self.config.values()]))

@@ -4,2452 +4,1517 @@
 ort 
-os%0Aimport yaml%0Aimport ujson as json%0Aimport rpy2%0Afrom rpy2 import robjects%0Afrom rpy2.robjects.packages import importr%0Afrom collections import namedtuple%0A%0Afrom survey_stats.log import logger%0Afrom survey_stats.survey import AnnotatedSurvey%0Afrom survey_stats.feathers import has_feather, load_feather%0A%0Arfther = importr('feather', on_conflict='warn')%0Arutils = importr('utils')%0Arbase = importr('base')%0A%0Acache_dir = os.path.join(os.getcwd(), 'cache')%0A%0A%0Aclass SurveyDataset(namedtuple('Dataset', %5B'config', 'surveys'%5D)):%0A    __slots__ = ()%0A%0A    @classmethod%0A    def load_dataset(cls, yml_f):%0A        config = None%0A        with open(yml_f, 'r') as fh:%0A            config = yaml.load(fh)%5B'surveys'%5D%0A        if not config:%0A            return cls(config=None, surveys=None)%0A        svys = %7B%7D%0A        for k, v in config.items():%0A            svys%5Bk%5D = cls.fetch_or_load_dataset(k, v%5B'spss'%5D, v%5B'data'%5D)%0A%0A        return cls(config=config, surveys=svys)%0A%0A    def fetch_or_load_dataset(id, spss_f, data_f):%0A        f = os.path.join(cache_dir, '%25s.feather' %25 id)%0A        ret = None%0A        rdf = None%0A        try:%0A            rdf = rfther.read_feather(f)%0A            ret = AnnotatedSurvey.from_rdf(spss_f, rdf)%0A            logger.info('loaded survey data from feather cache: %25s' %25 f)%0A        except:%0A            logger.info('could not find feather cache, loading raw data')%0A            ret = AnnotatedSurvey.load_cdc_survey(spss_f, data_f)%0A            logger.info('saving data to feather cache: %25s' %25 f)%0A            rfther.write_feather(ret.rdf, f)%0A%0A        rbase.gc()%0A        return ret%0A%0A%0Aclass YRBSSDataset(SurveyDataset):%0A    __slots__ = ()%0A%0A    def fetch_survey(self, combined=True, national=True, year=None):%0A        pred = (lambda v: v%5B'is_combined'%5D == combined and%0A                (v%5B'year'%5D == year if year else True))%0A        if not self.config:%0A            return None%0A        return next(( self.surveys%5Bk%5D for k, v in self.config.items() if pred(v)), None)%0A%0A    def fetch_config(self, national=True, year=None):%0A        combined = not year # if year is present, get ind year, else combined%0A        pred = (lambda v: v%5B'is_combined'%5D == combined and%0A                (v%5B'year'%5D == year if year else True))%0A        if not self.config:%0A            return None%0A        return next(((k, v) for k, v in self.config.items() if pred(v)), None)%0A%0A%0A    @property%0A    def survey_years(self):%0A        return set(sorted(%5Bint(v%5B'year'%5D) for v in self.config.values()%5D))
+blaze as bz%0Aimport feather%0Afrom collections import namedtuple%0A%0Afrom survey_stats import log%0Afrom survey_stats.survey import AnnotatedSurvey%0Afrom survey_stats.serdes import has_feather, load_feather%0A%0ASTATS_COLUMNS = %5B'count','mean','ci_u','ci_l','se'%5D%0A%0Alogger = log.getLogger()%0A%0Adef resolve_db_url(url):%0A    return (feather.read_dataframe(url) if%0A            url.endswith('.feather') else%0A            bz.data(url))%0A%0A%0A%0Aclass SurveyDataset(namedtuple('SurveyDataset',%0A                               %5B'svy_db','meta_db', 'cfg'%5D)):%0A%0A    @classmethod%0A    def from_db_urls(cls, svy_url, meta_url, cfg):%0A        return cls(svy_db=resolve_db_url(svy_url),%0A                   meta_db=resolve_db_url(meta_url), cfg=cfg)%0A%0A%0A    def fetch_socrata(self, qn, vars, filt = %7B%7D):%0A        sel = df.qid == qn%0A        if 'sitecode' in filt.keys():%0A            sel = sel & df.sitecode.isin(filt%5B'sitecode'%5D)%0A        if 'year' in filt.keys():%0A            sel = sel & df.year.isin(map( int, filt%5B'year'%5D))%0A        for v in self.facets.keys():%0A            if v in filt.keys():%0A                sel = sel & df%5Bv%5D.isin(filt%5Bv%5D)%0A            if not v in vars:%0A                sel = sel & (df%5Bv%5D == 'Total')%0A        cols = %5B'qid','response','sitecode','year'%5D + %5C%0A                self.facets + STATS_COLUMNS%0A        return df.to_dict(orient='records')%0A%0A%0A    def facets(self):%0A        return self.cfg%5B'facets'%5D%0A%0A%0A    def facet_levels(self):%0A        return %7Bk: list(self.meta_db%5Bk%5D%0A                        .cat.categories) for k in self.facets%7D%0A%0A
 %0A

update imports, remove unnecessary methods