<commit_before>import io
import json
import os
import tempfile
from typing import DefaultDict
import pandas as pd
from veneer.actions import get_big_data_source
import veneer
from string import Template
import argparse
import logging

from veneer.manage import VENEER_EXE_FN, create_command_line, kill_all_now
import veneer.manage as manage
logger = logging.getLogger(__name__)

def _BEFORE_BATCH_NOP(slf,x,y):
    pass

STANDARD_SOURCE_ELEMENTS=[
    'Downstream Flow Volume',
    'Upstream Flow Volume'
]

STANDARD_SOURCE_VARIABLES=[
    'Quick Flow',
    'Slow Flow',
    'Total Flow'
]
SOURCE_STORAGE_VARIABLES=[
    'Storage Volume',
    'Water Surface Area',
    'Water Surface Elevation',
    'Regulated Release Volume',
    'Spill Volume',
    'Total Inflow Volume',
    'Total Outflow Volume'
]

CONSTITUENT_VARIABLES=[
    'Downstream Flow Mass',
    'Total Flow Mass',
    'Stored Concentration',
    'Stored Mass'
]

CONSTITUENT_VARIABLE_LABELS=[
    'network',
    'generation',
    'storage_concentration',
    'storage_mass'
]

CONSTITUENT_NAME_FN=[
    veneer.name_for_location,
    veneer.name_for_fu_and_sc,
    veneer.name_for_location,
    veneer.name_for_location
]

class SourceExtractor(object):
    def __init__(self,v,dest,results=None,climate_data_sources=['Climate Data'],progress=logger.info):
        self.v=v
        self.dest=dest
        self.current_dest=None
        self.results = results
        if self.results is None:
            self.results = os.path.join(self.dest,'Results')
        self.climate_data_sources = climate_data_sources
        self.progress=progress

    def _ensure(self):
        if not os.path.exists(self.current_dest):
            os.makedirs(self.current_dest)

    def writeable(self,fn):
        self._ensure()
        return open(os.path.join(self.current_dest,fn),'w')

    def write_json(self,fn,data):
        json.dump(data,self.writeable(fn+'.json'))

    def write_csv(self,fn,df):
        self._ensure()
        fn = os.path.join(self.current_dest,fn+'.csv')+'.gz'
        df.to_csv(fn,compression='gzip')

    def _extract_structure(self):
        fus = set(self.v.model.catchment.get_functional_unit_types())
        constituents = self.v.model.get_constituents()
        constituent_sources = self.v.model.get_constituent_sources()
        assert len(constituent_sources)==1 # We don't use constituent source

        network = self.v.network()
        network_df = network.as_dataframe()

        fu_areas = self.v.retrieve_csv('/tables/fus')
        fu_areas = pd.read_csv(io.StringIO(fu_areas),index_col=0)

        self.write_json('constituents',constituents)
        self.write_json('fus',list(fus))
        self.writeable('network.json').write(network_df.to_json())
        self.write_csv('fu_areas',fu_areas)

    def _retrieve_input_timeseries(self,input_table,name_columns=['NetworkElement'],name_join=':'):
        result = pd.DataFrame()
        cols = set(input_table.columns) - set(name_columns)

        for _,row in input_table.iterrows():
            row_name = name_join.join([row[nc] for nc in name_columns])
            for col in cols:
                if pd.isna(row[col]) or not row[col]:
                    continue
                path = row[col]
                name = f'{row_name} {col}'
                path_components = path.split('/')
                data_source = path_components[2]
                input_set = path_components[3]
                item = path_components[4]
                time_series = self.v.data_source_item(data_source,item,input_set=input_set)
                result[name] = time_series[time_series.columns[0]]
        return result

    def _extract_runoff_configuration(self):
        self._extract_models_and_parameters('catchment.runoff','runoff_models','rr')

        self.progress('Getting climate data')
        climate = pd.DataFrame()
        for ds_name in self.climate_data_sources:
            ds = get_big_data_source(self.v,ds_name,self.data_sources,self.progress)
            climate = pd.concat([climate,ds],axis=1)

        self.write_csv('climate',climate)

    def _extract_models_and_parameters(self,path,model_fn,param_prefix,**kwargs):
        source = self.v.model
        for p in path.split('.'):
            source = getattr(source,p)
        models = source.model_table(**kwargs)
        params = source.tabulate_parameters(**kwargs)
        self.write_csv(model_fn,models)
        for model_type, table in params.items():
            self.write_csv('%s-%s'%(param_prefix,model_type),table)

    def _extract_piecewise_routing_tables(self):
        piecewise_links = self.v.model.link.routing.get_param_values('link.Name if not i_0.FlowRouting.IsGeneric else None')
        piecewise_links = [p for p in piecewise_links if p is not None]
        res = {}
        for pl in piecewise_links:
            flows = self.v.model.link.routing.get_param_values('Piecewises.*IndexFlow',links=pl)
            travel_times = self.v.model.link.routing.get_param_values('Piecewises.*TravelTime',links=pl)
            df = pd.DataFrame({'IndexFlow':flows,'TravelTime':travel_times})
            res[pl] = df
            if not len(df):
                self.progress(f'Expected piecewise routing table for link {pl}, but no rows')
                continue
            self.write_csv('piecewise-routing-%s'%(pl),df)

    def _extract_routing_configuration(self):
        self._extract_models_and_parameters('link.routing','routing_models','fr')
        self._extract_models_and_parameters('link.constituents','transportmodels','cr')

        self._extract_piecewise_routing_tables()
        # link_models = self.v.model.link.routing.model_table()

        # transport_models = self.v.model.link.constituents.model_table()
        # transport_params = self.v.model.link.constituents.tabulate_parameters()
        # self.write_csv('routing_models',link_models)
        # for model_type, table in transport_params.items():
        #     self.write_csv('fr-%s'%model_type,table)

        # # self.write_csv('routing_params',link_params)
        # self.write_csv('transportmodels',transport_models)

        # for model_type, table in transport_params.items():
        #     self.write_csv('cr-%s'%model_type,table)

    def _extract_generation_configuration(self):
        self._extract_models_and_parameters('catchment.generation','cgmodels','cg')
        # generation_models = self.v.model.catchment.generation.model_table()
        # generation_parameters = self.v.model.catchment.generation.tabulate_parameters()

        # for model_type, table in generation_parameters.items():
        #     self.write_csv('cg-%s'%model_type,table)

        # self.write_csv('cgmodels',generation_models)

    def _extract_demand_configuration(self):
        extraction_params = self.v.model.node.tabulate_parameters(node_types='ExtractionNodeModel')
        
        if not len(extraction_params):
            self.progress('No extraction points. Skipping water user extraction')
            return

        extraction_params = extraction_params['RiverSystem.Nodes.SupplyPoint.ExtractionNodeModel']
        # extractions = list(extraction_params['NetworkElement'])
        self.write_csv('extraction_point_params',extraction_params)

        water_users = self.v.model.node.water_users.names()
        if len(water_users):
            self.progress('Extracting information for %d water users'%len(water_users))
<commit_msg>Capture IsExtractive parameter from extraction nodes<commit_after>import io
import json
import os
import tempfile
from typing import DefaultDict
import pandas as pd
from veneer.actions import get_big_data_source
import veneer
from string import Template
import argparse
import logging

from veneer.manage import VENEER_EXE_FN, create_command_line, kill_all_now
import veneer.manage as manage
logger = logging.getLogger(__name__)

def _BEFORE_BATCH_NOP(slf,x,y):
    pass

STANDARD_SOURCE_ELEMENTS=[
    'Downstream Flow Volume',
    'Upstream Flow Volume'
]

STANDARD_SOURCE_VARIABLES=[
    'Quick Flow',
    'Slow Flow',
    'Total Flow'
]
SOURCE_STORAGE_VARIABLES=[
    'Storage Volume',
    'Water Surface Area',
    'Water Surface Elevation',
    'Regulated Release Volume',
    'Spill Volume',
    'Total Inflow Volume',
    'Total Outflow Volume'
]

CONSTITUENT_VARIABLES=[
    'Downstream Flow Mass',
    'Total Flow Mass',
    'Stored Concentration',
    'Stored Mass'
]

CONSTITUENT_VARIABLE_LABELS=[
    'network',
    'generation',
    'storage_concentration',
    'storage_mass'
]

CONSTITUENT_NAME_FN=[
    veneer.name_for_location,
    veneer.name_for_fu_and_sc,
    veneer.name_for_location,
    veneer.name_for_location
]

class SourceExtractor(object):
    def __init__(self,v,dest,results=None,climate_data_sources=['Climate Data'],progress=logger.info):
        self.v=v
        self.dest=dest
        self.current_dest=None
        self.results = results
        if self.results is None:
            self.results = os.path.join(self.dest,'Results')
        self.climate_data_sources = climate_data_sources
        self.progress=progress

    def _ensure(self):
        if not os.path.exists(self.current_dest):
            os.makedirs(self.current_dest)

    def writeable(self,fn):
        self._ensure()
        return open(os.path.join(self.current_dest,fn),'w')

    def write_json(self,fn,data):
        json.dump(data,self.writeable(fn+'.json'))

    def write_csv(self,fn,df):
        self._ensure()
        fn = os.path.join(self.current_dest,fn+'.csv')+'.gz'
        df.to_csv(fn,compression='gzip')

    def _extract_structure(self):
        fus = set(self.v.model.catchment.get_functional_unit_types())
        constituents = self.v.model.get_constituents()
        constituent_sources = self.v.model.get_constituent_sources()
        assert len(constituent_sources)==1 # We don't use constituent source

        network = self.v.network()
        network_df = network.as_dataframe()

        fu_areas = self.v.retrieve_csv('/tables/fus')
        fu_areas = pd.read_csv(io.StringIO(fu_areas),index_col=0)

        self.write_json('constituents',constituents)
        self.write_json('fus',list(fus))
        self.writeable('network.json').write(network_df.to_json())
        self.write_csv('fu_areas',fu_areas)

    def _retrieve_input_timeseries(self,input_table,name_columns=['NetworkElement'],name_join=':'):
        result = pd.DataFrame()
        cols = set(input_table.columns) - set(name_columns)

        for _,row in input_table.iterrows():
            row_name = name_join.join([row[nc] for nc in name_columns])
            for col in cols:
                if pd.isna(row[col]) or not row[col]:
                    continue
                path = row[col]
                name = f'{row_name} {col}'
                path_components = path.split('/')
                data_source = path_components[2]
                input_set = path_components[3]
                item = path_components[4]
                time_series = self.v.data_source_item(data_source,item,input_set=input_set)
                result[name] = time_series[time_series.columns[0]]
        return result

    def _extract_runoff_configuration(self):
        self._extract_models_and_parameters('catchment.runoff','runoff_models','rr')

        self.progress('Getting climate data')
        climate = pd.DataFrame()
        for ds_name in self.climate_data_sources:
            ds = get_big_data_source(self.v,ds_name,self.data_sources,self.progress)
            climate = pd.concat([climate,ds],axis=1)

        self.write_csv('climate',climate)

    def _extract_models_and_parameters(self,path,model_fn,param_prefix,**kwargs):
        source = self.v.model
        for p in path.split('.'):
            source = getattr(source,p)
        models = source.model_table(**kwargs)
        params = source.tabulate_parameters(**kwargs)
        self.write_csv(model_fn,models)
        for model_type, table in params.items():
            self.write_csv('%s-%s'%(param_prefix,model_type),table)

    def _extract_piecewise_routing_tables(self):
        piecewise_links = self.v.model.link.routing.get_param_values('link.Name if not i_0.FlowRouting.IsGeneric else None')
        piecewise_links = [p for p in piecewise_links if p is not None]
        res = {}
        for pl in piecewise_links:
            flows = self.v.model.link.routing.get_param_values('Piecewises.*IndexFlow',links=pl)
            travel_times = self.v.model.link.routing.get_param_values('Piecewises.*TravelTime',links=pl)
            df = pd.DataFrame({'IndexFlow':flows,'TravelTime':travel_times})
            res[pl] = df
            if not len(df):
                self.progress(f'Expected piecewise routing table for link {pl}, but no rows')
                continue
            self.write_csv('piecewise-routing-%s'%(pl),df)

    def _extract_routing_configuration(self):
        self._extract_models_and_parameters('link.routing','routing_models','fr')
        self._extract_models_and_parameters('link.constituents','transportmodels','cr')

        self._extract_piecewise_routing_tables()
        # link_models = self.v.model.link.routing.model_table()

        # transport_models = self.v.model.link.constituents.model_table()
        # transport_params = self.v.model.link.constituents.tabulate_parameters()
        # self.write_csv('routing_models',link_models)
        # for model_type, table in transport_params.items():
        #     self.write_csv('fr-%s'%model_type,table)

        # # self.write_csv('routing_params',link_params)
        # self.write_csv('transportmodels',transport_models)

        # for model_type, table in transport_params.items():
        #     self.write_csv('cr-%s'%model_type,table)

    def _extract_generation_configuration(self):
        self._extract_models_and_parameters('catchment.generation','cgmodels','cg')
        # generation_models = self.v.model.catchment.generation.model_table()
        # generation_parameters = self.v.model.catchment.generation.tabulate_parameters()

        # for model_type, table in generation_parameters.items():
        #     self.write_csv('cg-%s'%model_type,table)

        # self.write_csv('cgmodels',generation_models)

    def _extract_demand_configuration(self):
        constraint = dict(node_types='ExtractionNodeModel')
        extraction_params = self.v.model.node.tabulate_parameters(**constraint)
        
        if not len(extraction_params):
            self.progress('No extraction points. Skipping water user extraction')
            return

        extraction_params = extraction_params['RiverSystem.Nodes.SupplyPoint.ExtractionNodeModel']
        names = self.v.model.node.names(**constraint)
        is_extractive = self.v.model.node.get_param_values('IsExtractive',**constraint)
        extraction_params['IsExtractive'] = extraction_params.NetworkElement.map(dict(zip(names,is_extractive)))
        # extractions = list(extraction_params['NetworkElement'])
        self.write_csv('extraction_point_params',extraction_params)

        water_users = self.v.model.node.water_users.names()
        if len(water_users):
            self.progress('Extracting information for %d water users'%len(water_users))

{'fraction_non_alphanumeric': 0.11086299750426902, 'fraction_numerical': 0.0010508340995665309, 'mean_word_length': 3.1672140120415984, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '86240', 'n_tokens_mistral': 4794, 'n_tokens_neox': 4669, 'n_words': 749}