<commit_before># -*- coding:utf-8 -*-
# Copyright © 2011 <NAME>, <NAME>

"""
openFisca, Logiciel libre de simulation du système socio-fiscal français
Copyright © 2011 <NAME>, <NAME>

This file is part of openFisca.

    openFisca is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    openFisca is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with openFisca.  If not, see <http://www.gnu.org/licenses/>.
"""

from __future__ import division
import numpy as np
from pandas import DataFrame, read_csv, HDFStore
from src.lib.utils import of_import
from src.lib.description import ModelDescription, Description



class DataTable(object):
    """
    Construct a SystemSf object is a set of Prestation objects
        * title [string]
        * comment [string]: text shown on the top of the first data item
    """
    def __init__(self, model_description, survey_data = None, scenario = None, datesim = None, country = None):
        super(DataTable, self).__init__()

        # Init instance attribute
        self.description = None
        self.scenario = None
        self._isPopulated = False
        self.col_names = []
        self.table = DataFrame()
        self.index = {}
        self._nrows = 0
        
        if datesim is None:
            raise Exception('InputTable: datesim should be provided')
        else:
            self.datesim = datesim 
            
        if country is None:
            raise Exception('InputTable: country should be provided')
        else:
            self.country = country
                    
        self.survey_year = None
        
        # Build the description attribute        
        if type(model_description) == type(ModelDescription):
            descr = model_description()
            self.description = Description(descr.columns)
        else:
            raise Exception("model_description should be an ModelDescription inherited class")

        self.col_names = self.description.col_names

        if (survey_data is not None) and (scenario is not None):
            raise Exception("should provide either survey_data or scenario but not both")
        elif survey_data is not None:
            self.populate_from_survey_data(survey_data)
        elif scenario is not None:
            self.scenario = scenario
            scenario.populate_datatable(self)
        
    def gen_index(self, entities):
        '''
        Genrates indexex for the relevant entities
        '''
        
        self.index = {'ind': {0: {'idxIndi':np.arange(self._nrows), 
                                  'idxUnit':np.arange(self._nrows)}, # Units stand for entities
                      'nb': self._nrows},
                      'noi': {}}
        dct = self.index['noi']
        nois = self.table.noi.values
        listnoi = np.unique(nois)
        for noi in listnoi:
            idxIndi = np.sort(np.squeeze((np.argwhere(nois == noi))))
            idxUnit = np.searchsorted(listnoi, nois[idxIndi])
            temp = {'idxIndi':idxIndi, 'idxUnit':idxUnit}
            dct.update({noi: temp}) 

        for entity in entities:
            enum = self.description.get_col('qui'+entity).enum
            try:
                idx = getattr(self.table, 'id'+entity).values
                qui = getattr(self.table, 'qui'+entity).values
                enum = self.description.get_col('qui'+entity).enum
            except:
                raise Exception('DataTable needs columns %s and %s to build index with entity %s' %
                          ('id' + entity, 'qui' + entity, entity))

            self.index[entity] = {}
            dct = self.index[entity]
            idxlist = np.unique(idx)
            dct['nb'] = len(idxlist)

            for full, person in enum:
                idxIndi = np.sort(np.squeeze((np.argwhere(qui == person))))
#                if (person == 0) and (dct['nb'] != len(idxIndi)):
#                    raise Exception('Invalid index for %s: There is %i %s and %i %s' %(entity, dct['nb'], entity, len(idxIndi), full))
                idxUnit = np.searchsorted(idxlist, idx[idxIndi])
                temp = {'idxIndi':idxIndi, 'idxUnit':idxUnit}
                dct.update({person: temp}) 
    
    def propagate_to_members(self, entity , col):
        """
        Set the variable of all entity member to the value of the (head of) entity
        """
        index = self.index[entity]
        value = self.get_value(col, index)
        try:
            enum = self.description.get_col('qui'+entity).enum
        except:
            enum = self._inputs.description.get_col('qui'+entity).enum
        
        for member in enum:
            self.set_value(col, value, index, opt = member[1])


    def populate_from_survey_data(self, fname, year = None):
        '''
        Populates a DataTable from survey data
        '''

        if self.country is None:
            raise Exception('DataTable: country key word variable must be set') 
               
        # imports country specific variables from country.__init__.py
        INDEX = of_import(None, 'ENTITIES_INDEX', self.country)
        WEIGHT = of_import(None, 'WEIGHT', self.country)
        WEIGHT_INI = of_import(None, 'WEIGHT_INI', self.country)
        
        
        if fname[-4:] == '.csv':
            with open(fname) as survey_data_file:
                self.table = read_csv(survey_data_file)

        elif fname[-3:] == '.h5':
            store = HDFStore(fname)
            available_years = sorted([int(x[-4:]) for x in  store.keys()])
            
            if year is None:
                if self.datesim is not None:
                    year_ds  = self.datesim.year
                else:
                    raise Exception('self.datesim or year should be defined') 
            else:
                year_ds = year
           
            yr = year_ds+0 # to avoid pointers problem
            while yr not in available_years and yr > available_years[0]:
                yr = yr - 1
            base_name = 'survey_'+ str(yr)
            if year_ds != yr:
                print 'Survey data for year %s not found. Using year %s' %(str(year_ds), str(yr))
            else:
                print 'Survey data for year %s found' %str(year_ds)

            if yr in available_years:
                self.survey_year = yr
            self.table = store[str(base_name)] # only valid for frame 
            store.close()
            
        self._nrows = self.table.shape[0]
        missing_col = []
        for col in self.description.columns.itervalues():
            if not col.name in self.table:
                missing_col.append(col.name)
                self.table[col.name] = col._default
            try:   
                self.table[col.name] = self.table[col.name].astype(col._dtype)
            except:
                raise Exception("Impossible de lire la variable suivante issue des données d'enquête :\n %s \n  " %col.name)
            
        if missing_col:
            message = "%i input variables missing\n" % len(missing_col)
            for var in missing_col:
                message += '  - '+ var + '\n'
            print Warning(message)
        
        for var in INDEX:
            if ('id' + var) in missing_col:
                raise Exception('Survey data needs variable %s' % ('id' + var))
            
            if ('qui' + var) in missing_col:
                raise Exception('Survey data needs variable %s' % ('qui' + var))

        
        self.gen_index(INDEX)
        self._isPopulated = True
        
        # Initialize default weights
        self.set_value(WEIGHT_INI, self.get_value(WEIGHT),self.index['ind'])
        
#        # TODO: activate for debug
#        print self.table.get_dtype_counts()
#        
#        for col in self.table.columns:
#            if col not in self.description.col_names:
#                print 'removing : ',  col
#                del self.table[col]
#        
#        print self.table.get_dtype_counts()

    def get_value(self, varname, index = None, opt = None, sum_ = False):
        '''
        Read the value in an array
        
        Parameters
        ----------<commit_msg>better warning for missing col<commit_after># -*- coding:utf-8 -*-
# Copyright © 2011 <NAME>, <NAME>

"""
openFisca, Logiciel libre de simulation du système socio-fiscal français
Copyright © 2011 <NAME>, <NAME>

This file is part of openFisca.

    openFisca is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    openFisca is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with openFisca.  If not, see <http://www.gnu.org/licenses/>.
"""

from __future__ import division
import numpy as np
from pandas import DataFrame, read_csv, HDFStore
from src.lib.utils import of_import
from src.lib.description import ModelDescription, Description



class DataTable(object):
    """
    Construct a SystemSf object is a set of Prestation objects
        * title [string]
        * comment [string]: text shown on the top of the first data item
    """
    def __init__(self, model_description, survey_data = None, scenario = None, datesim = None, country = None):
        super(DataTable, self).__init__()

        # Init instance attribute
        self.description = None
        self.scenario = None
        self._isPopulated = False
        self.col_names = []
        self.table = DataFrame()
        self.index = {}
        self._nrows = 0
        
        if datesim is None:
            raise Exception('InputTable: datesim should be provided')
        else:
            self.datesim = datesim 
            
        if country is None:
            raise Exception('InputTable: country should be provided')
        else:
            self.country = country
                    
        self.survey_year = None
        
        # Build the description attribute        
        if type(model_description) == type(ModelDescription):
            descr = model_description()
            self.description = Description(descr.columns)
        else:
            raise Exception("model_description should be an ModelDescription inherited class")

        self.col_names = self.description.col_names

        if (survey_data is not None) and (scenario is not None):
            raise Exception("should provide either survey_data or scenario but not both")
        elif survey_data is not None:
            self.populate_from_survey_data(survey_data)
        elif scenario is not None:
            self.scenario = scenario
            scenario.populate_datatable(self)
        
    def gen_index(self, entities):
        '''
        Genrates indexex for the relevant entities
        '''
        
        self.index = {'ind': {0: {'idxIndi':np.arange(self._nrows), 
                                  'idxUnit':np.arange(self._nrows)}, # Units stand for entities
                      'nb': self._nrows},
                      'noi': {}}
        dct = self.index['noi']
        nois = self.table.noi.values
        listnoi = np.unique(nois)
        for noi in listnoi:
            idxIndi = np.sort(np.squeeze((np.argwhere(nois == noi))))
            idxUnit = np.searchsorted(listnoi, nois[idxIndi])
            temp = {'idxIndi':idxIndi, 'idxUnit':idxUnit}
            dct.update({noi: temp}) 

        for entity in entities:
            enum = self.description.get_col('qui'+entity).enum
            try:
                idx = getattr(self.table, 'id'+entity).values
                qui = getattr(self.table, 'qui'+entity).values
                enum = self.description.get_col('qui'+entity).enum
            except:
                raise Exception('DataTable needs columns %s and %s to build index with entity %s' %
                          ('id' + entity, 'qui' + entity, entity))

            self.index[entity] = {}
            dct = self.index[entity]
            idxlist = np.unique(idx)
            dct['nb'] = len(idxlist)

            for full, person in enum:
                idxIndi = np.sort(np.squeeze((np.argwhere(qui == person))))
#                if (person == 0) and (dct['nb'] != len(idxIndi)):
#                    raise Exception('Invalid index for %s: There is %i %s and %i %s' %(entity, dct['nb'], entity, len(idxIndi), full))
                idxUnit = np.searchsorted(idxlist, idx[idxIndi])
                temp = {'idxIndi':idxIndi, 'idxUnit':idxUnit}
                dct.update({person: temp}) 
    
    def propagate_to_members(self, entity , col):
        """
        Set the variable of all entity member to the value of the (head of) entity
        """
        index = self.index[entity]
        value = self.get_value(col, index)
        try:
            enum = self.description.get_col('qui'+entity).enum
        except:
            enum = self._inputs.description.get_col('qui'+entity).enum
        
        for member in enum:
            self.set_value(col, value, index, opt = member[1])


    def populate_from_survey_data(self, fname, year = None):
        '''
        Populates a DataTable from survey data
        '''

        if self.country is None:
            raise Exception('DataTable: country key word variable must be set') 
               
        # imports country specific variables from country.__init__.py
        INDEX = of_import(None, 'ENTITIES_INDEX', self.country)
        WEIGHT = of_import(None, 'WEIGHT', self.country)
        WEIGHT_INI = of_import(None, 'WEIGHT_INI', self.country)
        
        
        if fname[-4:] == '.csv':
            with open(fname) as survey_data_file:
                self.table = read_csv(survey_data_file)

        elif fname[-3:] == '.h5':
            store = HDFStore(fname)
            available_years = sorted([int(x[-4:]) for x in  store.keys()])
            
            if year is None:
                if self.datesim is not None:
                    year_ds  = self.datesim.year
                else:
                    raise Exception('self.datesim or year should be defined') 
            else:
                year_ds = year
           
            yr = year_ds+0 # to avoid pointers problem
            while yr not in available_years and yr > available_years[0]:
                yr = yr - 1
            base_name = 'survey_'+ str(yr)
            if year_ds != yr:
                print 'Survey data for year %s not found. Using year %s' %(str(year_ds), str(yr))
            else:
                print 'Survey data for year %s found' %str(year_ds)

            if yr in available_years:
                self.survey_year = yr
            self.table = store[str(base_name)] # only valid for frame 
            store.close()
            
        self._nrows = self.table.shape[0]
        missing_col = []
        for col in self.description.columns.itervalues():
            if not col.name in self.table:
                missing_col.append(col.name)
                self.table[col.name] = col._default
            try:   
                self.table[col.name] = self.table[col.name].astype(col._dtype)
            except:
                raise Exception("Impossible de lire la variable suivante issue des données d'enquête :\n %s \n  " %col.name)
            
        if missing_col:
            message = "%i input variables missing\n" % len(missing_col)
            messagef = ""
            messageb = ""
            missing_col.sort()
            for var in missing_col:
                if var[0] == 'f':
                    messagef += '  - '+ var +'\n'
                elif var[0] == 'b':
                    messageb += '  - '+ var +'\n'
                else:
                    message += '  - '+ var +'\n'
            print Warning(message + messagef + messageb)
        
        for var in INDEX:
            if ('id' + var) in missing_col:
                raise Exception('Survey data needs variable %s' % ('id' + var))
            
            if ('qui' + var) in missing_col:
                raise Exception('Survey data needs variable %s' % ('qui' + var))

        
        self.gen_index(INDEX)
        self._isPopulated = True
        
        # Initialize default weights
        self.set_value(WEIGHT_INI, self.get_value(WEIGHT),self.index['ind'])
        
#        # TODO: activate for debug
#        print self.table.get_dtype_counts()
#        
#        for col in self.table.columns:
#            if col not in self.description.col_names:
#                print 'removing : ',  col
#                del self.table[col]
#        
#        print self.table.get_dtype_counts()

    def get_value(self, varname, index = None, opt = None, sum_ = False):
        '''
        Read the value in an array
        
        Parameters
        ----------
{'fraction_non_alphanumeric': 0.09787015360018272, 'fraction_numerical': 0.002626620224975732, 'mean_word_length': 1.7822081016679905, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3908411', 'n_tokens_mistral': 4959, 'n_tokens_neox': 4754, 'n_words': 1450}