text:
# encoding: utf-8

'''
@author: Majeed.sahebzadha
'''

import numpy as np
import pandas as pd
from math import ceil
import re
import operator
from quantipy.core.helpers import functions as helpers

''' Simplified access to, and manipulation of, the pandas dataframe.
    Contains various helper functions.
'''

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def clean_df_values(old_df, replace_this, replace_with_that, regex_bol, as_type):
    '''

    '''

    new_df = old_df.replace(replace_this,
                            replace_with_that,
                            regex=regex_bol).astype(as_type)

    return new_df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def case_insensitive_matcher(check_these, against_this):
    '''
    performs the case-insensitive search of given list of items against df rows or columns and pulls out
    matched items from the df.
    '''
    matched = [v
               for x,d in enumerate(check_these)
                   for i,v in enumerate(against_this)
                        if v.lower() == d.lower()
                        ]
    return matched

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def remove_percentage_sign(old_df):
    '''

    '''

    new_df = old_df.replace('%','',regex=True).astype(np.float64)

    return new_df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def drop_null_rows(old_df, axis_type=1):
    '''
    drop rows with all columns having value 0
    '''

    new_df = old_df.loc[(df!=0).any(axis=axis_type)]

    return new_df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def auto_sort(df, fixed_categories=[], column_position=0, ascend=True):
    '''
    Sorts a flattened (non multiindexed) panda dataframe whilst excluding given rows

    Params:
    -------

    df: pandas dataframe
    fixed_categories: list
        list of row labels
        example: [u'Other', u"Don't know/ can't recall", u'Not applicable']
    column_position: int
    ascend: boolean
        Sort ascending vs. descending
    '''

    # ensure df is not empty
    if not df.empty:

        # ensure df is not multiindexed
        nblevels = df.index.nlevels
        if nblevels == 2:
            raise Exception('Expected Flat DF, got multiindex DF')

        # ensure fixed_categories is not empty
        if fixed_categories:

            #reindex df because it might contain duplicates
            df = df.reset_index()

            #df with no fixed categories, then sort.
            df_without_fc = df.loc[~df[df.columns[0]].isin(fixed_categories)]
            if pd.__version__ == '0.19.2':
                df_without_fc = df_without_fc.sort_values(by=df.columns[column_position+1], ascending=ascend)
            else:
                df_without_fc = df_without_fc.sort(columns=df.columns[column_position+1], ascending=ascend)

            #put each row as a tuple in a list
            tups = []
            for x in df_without_fc.itertuples():
                tups.append(x)

            #get fixed categories as a df
            df_fc = df[~df.index.isin(df_without_fc.index)]

            #convert fixed categories to rows of tuples,
            #then insert row to tups list in a specific index
            for x in df_fc.itertuples():
                tups.insert(x[0], x)

            #remove the indexes from the list of tuples
            filtered_tups = [x[1:] for x in tups]

            #put all the items in the tups list together to build a df
            new_df = pd.DataFrame(filtered_tups, columns=list(df.columns.values))
            new_df = new_df.set_index(df.columns[0])

        else:
            if pd.__version__ == '0.19.2':
                new_df = df.sort_values(by=df.columns[column_position], ascending=ascend)
            else:
                new_df = df.sort(columns=df.columns[column_position], ascending=ascend)

    return new_df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def all_same(numpy_list):
    '''

    '''

    val = numpy_list.tolist()

    return all(x == val[0] for x in val)

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def find_dups(df, orientation='Side'):
    '''
    Looks for duplicate labels in a df. Converts axis
    labels to a list and then returns duplicate index from list.
    If the list contains duplicates then a statememnt is returned.
    '''

    if orientation == 'Side':
        mylist = list(df.index.values)
        axis = 'row'
    else:
        mylist = list(df.columns.values)
        axis = 'column'

    dup_idx = [i for i, x in enumerate(mylist) if mylist.count(x) > 1]

    if dup_idx:
        statement = ("\n{indent:>10}*Warning: This table/chart contains duplicate "
                    "{orientation} labels".format(
                                             indent='',
                                             orientation=axis))
    else:
        statement = ''

    return statement

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def df_splitter(df, min_rows, max_rows):
    '''
    returns a list of dataframes sliced as evenly as possible
    '''

    #ensure the indexs are strings not ints or floats
    if not isinstance(df.index, str):
        df.index = df.index.map(str)

    row_count = len(df.index)

    maxs = pd.Series(list(range(min_rows, max_rows+1)))
    rows = pd.Series([row_count]*maxs.size)
    mods = rows % maxs
    splitter = maxs[mods >= min_rows].max()

    if row_count <= max_rows:
        splitter = 1
    else:
        splitter = ceil(row_count/float(splitter))

    size = int(ceil(float(len(df)) / splitter))

    return [df[i:i + size] for i in range(0, len(df), size)]

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def strip_html_tags(text):
    '''
    Strip HTML tags from any string and transform special entities
    '''

    rules = [
             {r'<[^<]+?>': ''},                # remove remaining tags
             {r'^\s+' : '' },                  # remove spaces at the beginning
             {r'\,([a-zA-Z])': r', \1'},        # add space after a comma
             {r'\s+' : ' '}                    # replace consecutive spaces
             ]
    for rule in rules:
        for (k,v) in list(rule.items()):
            regex = re.compile(k)
            text = regex.sub(v, text)

    # replace special strings
    special = {
               ' ': ' ',
               '&': '&',
               '"': '"',
               '<': '<',
               '>': '>',
               '**': '',
               "’": "'"

               }
    for (k,v) in list(special.items()):
        text = text.replace(k, v)

    return text

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def clean_axes_labels(df):
    '''
    Cleans dataframe labels. Strips html code, double white spaces and so on.

    Params:
    -------

    df: pandas dataframe
    '''

    #standardise all index/column elements as unicode
    df_index_labels = df.index.map(str)
    df_col_labels = df.columns.map(str)

#     df_index_labels = [unicode(w)
#                        if not isinstance(w, unicode) and not isinstance(w, str)
#                        else w
#                        for w in df.index.values]

#     df_col_labels = df.columns.values

    col_labels = []
    index_labels = []

    for ctext in df_col_labels:
        ctext = strip_html_tags(ctext)
        col_labels.append(ctext)

    for indtext in df_index_labels:
        indtext = strip_html_tags(indtext)
        index_labels.append(indtext)

    df.columns = col_labels
    df.index = index_labels

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def color_setter(numofseries, color_order='reverse'):
    '''

    '''

    color_set = [(147,208,35), (83,172,175), (211,151,91), (17,124,198), (222,231,5), (136,87,136),
                 (88,125,21), (49,104,106), (143,91,38), (10,74,119), (133,139,3), (82,52,82),
                 (171,224,72), (117,189,191), (220,172,124), (38,155,236), (242,250,40), (165,115,165),
                 (118,166,28), (66,138,141), (190,121,51), (14,99,158), (178,185,4), (109,70,109),
                 (192,232,118), (152,205,207), (229,193,157), (92, 180,241), (245,252,94), (188,150,188),
                 (74,104,18), (41,87,88), (119,75,32), (9,62,99), (111,116,3), (68,44,68),
                 (197,227,168), (176,209,210), (229,199,178), (166,188,222), (235,240,166), (193,177,193),
                 (202,229,175), (182,212,213), (231,203,184), (174,194,224), (237,241,174), (198,183,198)]

    color_set = color_set[0:numofseries]
    if color_order == 'reverse':
        return color_set[::-1]
    else:
        return color_set

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def place_vals_in_labels(old_df, base_position=0, orientation='side', drop_position=True):
    '''
    Takes values from a given column or row and inserts it to the df's row or column labels.
    Normally used to insert base values in row or column labels.
    '''

    if orientation == 'side':
        #grab desired column's values, normally index 0
        col_vals = old_df.ix[:,[base_position]].values
        #col_vals returns a list of list which needs flattening
        flatten_col_vals = [item for sublist in col_vals for item in sublist]
        #grab row labels
        index_labels = old_df.index

        new_labels_list = {}
        for x,y in zip(index_labels, flatten_col_vals):
            new_labels_list.update({x : x + " (n=" + str(y) +")"})

        new_df = old_df.rename(index=new_labels_list, inplace=False)

        if drop_position:
            new_df = new_df.drop(new_df.columns[[base_position]], axis=1, inplace=False)

    else:
        #grab desired row's values, normally index 0
        row_vals = old_df.ix[[base_position],:].values
        #row_vals returns a list of list which needs flattening
        flatten_col_vals = [item for sublist in row_vals for item in sublist]
        #grab row labels
        col_labels = df.columns

        #rename rows one by one.
        new_labels_list = {}
        for x,y in zip(index_labels, flatten_col_vals):
            new_labels_list.update({x : x + " (n=" + str(y) +")"})

        new_df = old_df.rename(columns=new_labels_list, inplace=False)

        if drop_position:
            new_df = new_df.drop(new_df.index[[base_position]], axis=0, inplace=False)

    return new_df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def get_qestion_labels(cluster_name, meta, table_name=None):
    '''

    '''

    question_label_dict ={}

    text_key = meta['lib']['default text']

    table_list = list(cluster_name.keys())

    for table in table_list:
        view = cluster_name[table][cluster_name[table].data_key][cluster_name[table].filter][table][cluster_name[table].content_of_axis[0]][cluster_name[table].views[1]]

        vdf = view.dataframe
#         vdf = drop_hidden_codes(vdf, meta)
#         df = index_from_meta(vdf, meta, vdf)
#         question_label = df.index.get_level_values(level=0)[0]

#         question_label_dict[table] = question_label

        qname = vdf.index.get_level_values(0).tolist()[0]
        vdf_meta = meta['columns'].get(qname, '%s not in the columns set in the meta' % (qname))

        question_label_dict[table] = vdf_meta['text'][text_key]

#         question_label_dict[table] = vdf_meta['text'][text_key]

    if table_name:
        return question_label_dict[table_name]
    else:
        return question_label_dict

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def validate_cluster_orientations(cluster):
    '''
    Make sure that the chains follow the rule:
        - All chains must have the same orientation, x or y.
    '''
    if len(set([
        cluster[chain_name].orientation
        for chain_name in list(cluster.keys())
    ])) != 1:
        raise Exception(
            "Chain orientations must be consistent. Please review chain "
            "specification"
        )

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def get_base(df, base_description, is_mask):
    '''
    Constructs base text for any kind of chart, single, multiple, grid.

    Params:
    -------

    df: pandas dataframe
    base_description: str
    '''
    num_to_str = lambda string: str(int(round(string)))
    base_text_format = lambda txt, num: '{} ({})'.format(txt, num_to_str(num))

    #standardise all index/column elements as unicode
    df_index_labels = df.index.map(str)
    df_col_labels = df.columns.map(str)

    # get col labels and row values
    top_members = df.columns.values
    base_values = df.values

    # count row/col
    numofcols = len(df.columns)
    numofrows = len(df.index)

    #if base description is empty then
    if base_description:
        #example of what the base description would look like - 'Base: Har minst ett plagg'
        #remove the word "Base:" from the description
        description = base_description.split(': ')[-1]
        #grab the label for base from the df
        base_label = df.index[0]
        #put them together
        base_description = '{}: {}'.format(base_label, description)
    else:
        base_description = df.index.values[0]
    base_description = base_description.strip()

    #single series format
    if numofcols == 1:
        base_text = base_text_format(base_description, base_values[0][0])

    #multi series format
    elif numofcols > 1:
        # if all_same(base_values[0]):
        #     base_text = base_text_format(base_description, base_values[0][0])
        # else:
        if not is_mask:
            it = list(zip(top_members, base_values[0]))
            base_texts = ', '.join([base_text_format(x, y) for x, y in it])
            base_text = ' - '.join([base_description, base_texts])
        else:
            base_text = base_text_format(base_description, base_values[0][0])
    return base_text

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def replace_decimal_point_with(df, replacer=","):
    '''

    '''

    for col in df.columns:
        df[col] = pd.Series(["{0}".format(val) for val in df[col]], index = df.index)
    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def set_column_sequence(dataframe, seq):
    '''
    Takes a dataframe and a subsequence of its columns, returns dataframe with seq as first columns
    '''

    cols = seq[:] # copy so we don't mutate seq
    for x in dataframe.columns:
        if x not in cols:
            cols.append(x)

    return dataframe[cols]

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def round_df_cells(df, decimal_points):
    '''

    '''

    if decimal_points == 0:
        df = df.applymap(lambda x: int(round(x)))
    else:
        df = df.applymap(lambda x: round(x, decimal_points))

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def reverse_order(old_df, orientation='side'):
    '''
    Will reverse the order of rows or columns
    '''

    if orientation == 'side':
        df = old_df.iloc[::-1]
    else:
        df = old_df[old_df.columns[::-1]]

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def get_selection_by_index(df, position, orientation='side'):
    '''
    Grabs and returns a single column or row.
    example: myrow = get_selection_by_index(mydf, 2, 'side')
    grabs row 2 from df.
    '''

    if orientation == 'side':
        #will return a single row
        df = df.ix[[position],:]
    else:
        #will return a single col
        df = df.ix[:,[position]]

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def del_by_label(df, label_to_del, orientation='side'):
    '''
    deletes a single or multiple row or col labels from df
    param - label_to_del: takes a list of labels
    '''

    #if what's passed into label_to_del is not in a list then
    #put it in a list
    if not isinstance(label_to_del, list):
        label_to_del = [label_to_del]

    if orientation=='side':
        orientation=0
    else:
        orientation=1

    df = df.drop([label_to_del], axis=orientation, inplace=True)

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def rename_label(df, old_label, new_label, orientation='side'):
    '''
    Renames a single row or cols label
    '''

    if orientation == 'side':
        df.rename(index={old_label: new_label}, inplace=True)
    else:
        df.rename(columns={old_label: new_label}, inplace=True)

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def drop_hidden_codes(view):
    '''

    '''

    #drop hidden codes
    if 'x_hidden_codes' in view.meta():
        vdf = helpers.deep_drop(
                view.dataframe,
                view.meta()['x_hidden_codes'],
                axes=0)
    else:
        vdf = view.dataframe

    return vdf

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def paint_df(vdf, view, meta, text_key):
    '''

    '''

    #add question and value labels to df
    if 'x_new_order' in view.meta():
        df = helpers.paint_dataframe(
                df=vdf.copy(),
                meta=meta,
                ridx=view.meta()['x_new_order'],
                text_key=text_key)
    else:
        df = helpers.paint_dataframe(
                df=vdf.copy(),
                meta=meta,
                text_key=text_key)

    return df

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def partition_view_df(view, values=False, data_only=False, axes_only=False):
    '''
    Disassembles a view dataframe object into its
    inner-most index/columns parts (by dropping the first level)
    and the actual data.

    Parameters
    ----------
    view : Quantipy view

    values : boolean, optional
        If True will return the np.array
        containing the df values instead of a dataframe

    data_only : boolean, optional
        If True will only return the data component of the view dataframe

    axes_only : boolean, optional
        If True will only return the inner-most index and columns component
        of the view dataframe.

    Returns
    -------
    data, index, columns : dataframe (or np.array of values), index, columns
    '''
    df = view.copy()
    if isinstance(df.index, pd.MultiIndex):
        df.index = df.index.droplevel()
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.droplevel()
    index = df.index
    columns = df.columns
    data = df if not values else df.values

    if data_only:
        return data
    elif axes_only:
        return index.tolist(), columns.tolist()
    else:
        return data, index.tolist(), columns.tolist()

'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
'~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'

def is_grid_element(table_name, table_pattern):
    '''
    Checks if a table is a grid element or not

    Parameters
    ----------

    '''

    matches = table_pattern.findall(table_name)

    if (len(matches)>0 and len(matches[0])==2):
        matched = True
    else:
        matched = False

    return matched


text:
"""Interface for loading phystables trials"""
from ..space import WallSpace
from ..goals import BoxGoal2D
import warnings
import numpy as np

def load_phystable_trial(trial, goaltypes=[], ball_expand=True):
    """Loads in a phystables trial to use RRT on

    Currently outputs a WallSpace... may need to update for odd walls

    Args:
        trial (SimpleTrial): a trial from phystables for loading
        goaltypes ([returns]): a set of onreturns for the goals to turn into
                               goals for RRT
        ball_expand (bool): Should walls be expanded to account for ball size?

    Returns:
        A WallSpace analog of the trial
    """
    try:
        from phystables import SimpleTrial
    except:
        raise ImportError("phystables package is required for loading")

    assert isinstance(trial, SimpleTrial), "Must load a phystables trial"

    if len(trial.abnormwalls) > 0:
        warnings.warn("Abnormal walls will be ignored (not yet implemented)")

    # Go through the walls & goals to parse into Space walls & goals
    glist = []
    wlist = []
    if ball_expand:
        expd = trial.ball[2]
    else:
        expd = 0

    for w in trial.normwalls:
        wlist.append([np.array(w[0]) - expd, np.array(w[1]) + expd])
    for g in trial.goals:
        if g[2] in goaltypes:
            glist.append(BoxGoal2D(np.array(g[0]) - expd,
                                   np.array(g[1]) + expd))
        else:
            # Unloved goals become walls
            wlist.append([np.array(g[0]) - expd,
                          np.array(g[1]) + expd])

    return WallSpace(trial.dims, glist, wlist)

ext: py

max_stars_repo_name: kasmith/rrt

max_issues_repo_name: kasmith/rrt

max_forks_repo_name: kasmith/rrt

reason: import numpy


text:
import numpy as np
import pandas as pd
import os
import json
import typing as tp
import time
import math
import requests

from concurrent.futures import ThreadPoolExecutor

from brownie import network, Contract
from datetime import datetime, timedelta

from influxdb_client import InfluxDBClient
from influxdb_client.client.write_api import SYNCHRONOUS, PointSettings

BLOCK_SUBGRAPH_ENDPOINT = "https://api.thegraph.com/subgraphs/name/decentraland/blocks-ethereum-mainnet"  # noqa

def get_b_q(timestamp: int) -> str:
    return """query {
                blocks(
                    first: 1,
                    orderBy: timestamp,
                    orderDirection: desc,
                    where: { timestamp_lt: %s }
                ) {
                    timestamp
                    number
                }
            } """ % (str(timestamp))

def get_config() -> tp.Dict:
    return {
        "token": os.getenv('INFLUXDB_TOKEN'),
        "org": os.getenv('INFLUXDB_ORG'),
        "bucket": os.getenv('INFLUXDB_BUCKET', 'ovl_univ3_1m'),
        "url": os.getenv("INFLUXDB_URL"),
        "window": os.getenv("WINDOW", 60)
    }

def create_client(config: tp.Dict) -> InfluxDBClient:
    return InfluxDBClient(
        url=config['url'],
        token=config['token'],
        debug=False
    )

def get_point_settings() -> PointSettings:
    point_settings = PointSettings(**{"type": "metrics-minute"})
    point_settings.add_default_tag("influx-univ3", "ingest-data-frame")
    return point_settings

def POOL(addr: str, abi) -> Contract:
    # UniswapV3Pool contract
    return Contract.from_abi("unipool", addr, abi)

def get_quote_path() -> str:
    base = os.path.dirname(os.path.abspath(__file__))
    qp = 'constants/univ3_quotes.json'
    return os.path.join(base, qp)

def get_quotes() -> tp.List:
    quotes = []
    p = get_quote_path()
    with open(p) as f:
        data = json.load(f)
        quotes = data.get('quotes', [])
    return quotes

def get_uni_abi_path() -> str:
    base = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(base, 'constants/univ3.abi.json')

def get_uni_abi() -> tp.Dict:
    with open(get_uni_abi_path()) as f:
        data = json.load(f)
    return data

def get_b_t(timestamp: int) -> tp.Tuple:

    q = {'query': get_b_q(timestamp)}

    retries = 1
    success = False

    while not success and retries < 5:
        try:
            result = json.loads(
                requests.post(BLOCK_SUBGRAPH_ENDPOINT, json=q).text
            )['data']['blocks'][0]
            success = True
        except Exception as e:
            wait = retries * 10
            err_cls = e.__class__
            err_msg = str(e)
            msg = f'''
            Error type = {err_cls}
            Error message = {err_msg}
            Wait {wait} secs
            '''
            print(msg)
            time.sleep(wait)
            retries += 1

    return (int(result['number']), int(result['timestamp']))

def get_calls(
        pair: Contract,
        t_from: int,
        t_to: int,
        t_period: int) -> tp.List:

    calls = []

    while t_from < t_to:
        calls.append((pair, t_from, t_period))
        t_from += t_period

    return calls

def find_start(api, quote, config) -> int:

    retries = 1
    success = False

    while not success and retries < 5:
        try:
            r = api.query_data_frame(org=config['org'], query=f'''
                from(bucket:"{config['bucket']}")
                    |> range(start:0, stop: now())
                    |> filter(fn: (r) => r["id"] == "{quote['id']}")
                    |> filter(fn: (r) => r["_field"] == "tick_cumulative")
                    |> last()
            ''')
            success = True
        except Exception as e:
            wait = retries * 10
            err_cls = e.__class__
            err_msg = str(e)
            msg = f'''
            Error type = {err_cls}
            Error message = {err_msg}
            Wait {wait} secs
            '''
            print(msg)
            time.sleep(wait)
            retries += 1

    if (len(r.index) > 0):
        return math.floor(datetime.timestamp(
            r.iloc[0]['_time']))
    else:
        # return quote['time_deployed'] + 1200
        return math.floor(
            int(datetime.timestamp(datetime.now() - timedelta(days=60)))
            )

def read_cumulatives(args: tp.Tuple) -> tp.Tuple:

    (pair, t_at) = args
    (b, b_t) = get_b_t(t_at)
    (cum_tick, cum_liq) = pair.observe([0], block_identifier=b)
    return (b_t, cum_tick[0], cum_liq[0])

def list_cumulatives(args: tp.Tuple) -> tp.Tuple:

    (quote, pool, ts) = args
    item = read_cumulatives((pool, ts))
    print("item", item)
    print('time', datetime.fromtimestamp(
                item[0]).strftime("%m/%d/%Y, %H:%M:%S"))
    print('quote', quote['id'])

    return (
        datetime.utcfromtimestamp(float(item[0])),
        float(item[1]),
        float(item[2]),
        quote['token0_name'],
        quote['token1_name'],
        quote['id']
        )

def write_cumulatives(config, vals):

    df = pd.DataFrame(
            vals,
            columns=[
                '_time',
                'tick_cumulative',
                'liquidity_cumulative',
                'token0_name',
                'token1_name',
                'id'
                ]
            )
    df = df.set_index("_time")

    with InfluxDBClient(
            url=config['url'],
            token=config['token'],
            org=config['org']
            ) as client:

        print("Start ingestion")

        write_api = client.write_api(
            write_options=SYNCHRONOUS, point_settings=get_point_settings())

        j = 1
        while j == 1:
            try:
                write_api.write(bucket=config['bucket'],
                                record=df,
                                data_frame_measurement_name="mem",
                                data_frame_tag_columns=[
                                    'id', 'token0_name', 'token1_name'
                                    ]
                                )
                j = 0
                print("Ingested to influxdb")
            except Exception as e:
                err_cls = e.__class__
                err_msg = str(e)
                msg = f'''
                Error type = {err_cls}
                Error message = {err_msg}
                Wait 5 secs
                '''
                print(msg)
                continue

def get_uni_cumulatives(quotes, query_api, config, t_end):
    abi = get_uni_abi()
    t_step = 500

    for q in quotes:

        pool = POOL(q['pair'], abi)
        batch_size = (t_step * config['window'])
        t_start = find_start(query_api, q, config)
        t_interm = t_start + batch_size
        if t_interm > t_end:
            t_interm = t_end

        while t_start < t_end:
            list_cumulatives_calls = [
                (q, pool, x)
                for x in np.arange(
                    t_start,
                    t_interm,
                    config['window']
                    )
                ]

            list_cumulatives_vals = []
            with ThreadPoolExecutor() as executor:
                for item in executor.map(
                                        list_cumulatives,
                                        list_cumulatives_calls):
                    list_cumulatives_vals.append(item)

            write_cumulatives(config, list_cumulatives_vals)
            t_start = t_interm
            t_interm = t_start + batch_size
            if t_interm > t_end:
                t_interm = t_end

def main():
    print(f"You are using the '{network.show_active()}' network")

    config = get_config()
    quotes = get_quotes()
    client = create_client(config)
    query_api = client.query_api()

    t_end = math.floor(time.time())
    while True:
        get_uni_cumulatives(quotes, query_api, config, t_end)
        sl_time = config['window']
        print(f'Wait {sl_time} secs')
        time.sleep(sl_time)
        t_end = math.floor(time.time())

max_stars_repo_name: atitawat-pol/overlay-risk

max_issues_repo_name: atitawat-pol/overlay-risk

max_forks_repo_name: atitawat-pol/overlay-risk


text:
################################################################################
# A naive bayes classifier for classifying iris data.
#
# Author: Carl Cortright
# Date: 9/10/2016
# Updated: 9/12/2016
#
# Copyright 2016 Carl Cortright
################################################################################
import csv
import math
import random
import matplotlib.pyplot as plt
import numpy as np

#
# Returns a dictionary of probabilities of a feature given a class in the form:
# {class:{(position, value): probability ... }...}
#
def getIrisModel(inFile):
    # Set up the CSV file
    iris_data = open(inFile, "rt")
    iris_data_csv = csv.reader(iris_data, delimiter=',')
    # Read the CSV file into a 3d dict
    iris_dict = {"Iris-setosa":{}, "Iris-versicolor":{}, "Iris-virginica":{}}
    for i, row in enumerate(iris_data_csv):
        if(row[4] == "Iris-setosa"):
            iris_dict["Iris-setosa"][i] = row[0:4]
        elif(row[4] == "Iris-virginica"):
            iris_dict["Iris-virginica"][i] = row[0:4]
        elif(row[4] == "Iris-versicolor"):
            iris_dict["Iris-versicolor"][i] = row[0:4]

    iris_data.close()

    # Calculate all of the probabilities we need
    feature_probabilities = {"Iris-setosa":{}, "Iris-versicolor":{}, "Iris-virginica":{}}
    # Store the overall probability of a class
    NUM_ENTRIES = NUM_ENTRIES = len(iris_dict["Iris-setosa"]) + len(iris_dict["Iris-versicolor"]) + len(iris_dict["Iris-virginica"])
    feature_probabilities["Iris-setosa-p"] = float(len(iris_dict["Iris-setosa"]))/NUM_ENTRIES
    feature_probabilities["Iris-virginica-p"] = float(len(iris_dict["Iris-virginica"]))/NUM_ENTRIES
    feature_probabilities["Iris-versicolor-p"] = float(len(iris_dict["Iris-versicolor"]))/NUM_ENTRIES

    # Compute the conditional probabilities of feature x_i given C
    # Each entry is in the form of (position, value): probability
    for classification in iris_dict:
        for entry in iris_dict[classification].values():
            for i, value in enumerate(entry):
                if(not (i,value) in feature_probabilities[classification]):
                    feature_count = 0
                    for entry in iris_dict[classification].values():
                        if(entry[i] == value):
                            feature_count += 1
                    feature_probabilities[classification][(i,value)] = float(feature_count)/50

    return feature_probabilities

#
# Classify a given vector v based on the model.
# Returns the predicted classification.
#
def classify(v, model):
    # Probabilities of each class
    setosa_p = 0
    versicolor_p = 0
    virginica_p = 0
    # Probabilities of each class across the whole dataset
    IRIS_SETOSA_P = model["Iris-setosa-p"]
    IRIS_VERSICOLOR_P = model["Iris-versicolor-p"]
    IRIS_VIRGINICA_P = model["Iris-virginica-p"]
    for i, f_i in enumerate(v):
        # Using log-space
        if((i, f_i) in model["Iris-setosa"]):
            setosa_p += math.log(model["Iris-setosa"][(i, f_i)] + 1)
        if((i, f_i) in model["Iris-versicolor"]):
            versicolor_p += math.log(model["Iris-versicolor"][(i, f_i)] + 1)
        if((i, f_i) in model["Iris-virginica"]):
            virginica_p += math.log(model["Iris-virginica"][(i, f_i)] + 1)
    setosa_p *= IRIS_SETOSA_P
    versicolor_p *= IRIS_VERSICOLOR_P
    virginica_p *= IRIS_VIRGINICA_P
    probabilities = {setosa_p:"Iris-setosa", versicolor_p:"Iris-versicolor", virginica_p:"Iris-virginica"}
    return probabilities[max(probabilities)]

#
# Calculates the accuracy of the model
#
def calcAccuracy(training_data_file, test_data_file):
    naive_bayes_model = getIrisModel(training_data_file)
    test_data = open(test_data_file, "rt")
    test_data_csv = csv.reader(test_data, delimiter=",")
    test_data_dict = {}
    for i, row in enumerate(test_data_csv):
        test_vector = []
        for j in range(0,5):
            test_vector.append(row[j])
        test_data_dict[i] = test_vector

    correct_count = 0
    for entry in test_data_dict.values():
        if(classify(entry, naive_bayes_model) == entry[4]):
            correct_count += 1
    accuracy = float(correct_count)/len(test_data_dict)
    test_data.close()
    return accuracy

#
# Shuffles the dataset based on a ratio training:test
#
def shuffleData(training, test):
    iris_data = open("iris.data", "r+")
    iris_training = open("iris.training.data", "w")
    iris_test = open("iris.test.data", "w")

    all_data = iris_data.readlines()
    random.shuffle(all_data)
    all_data_len = len(all_data)

    split_point = int( all_data_len * ( float(training) / (training + test)))
    for i in range(0, split_point):
        iris_training.write(all_data[i])
    for j in range(split_point, all_data_len - 1):
        iris_test.write(all_data[j])

    iris_data.close()
    iris_training.close()
    iris_test.close()

#
# Calculates the accuracy distribution given an iris data file, returns dict
#
def calcAccuracyDist():
    ratio_max = 15
    accuracyDist = {}
    for ratio in range(1, ratio_max):
        total = 0
        tests = 20000
        for trial in range(0, tests):
            shuffleData(ratio, 1)
            total += calcAccuracy("iris.training.data", "iris.test.data")
        average = total/tests
        average = average * 100
        print("(" + str(ratio) + ":1) = " + str(average) + "%")
        accuracyDist[ratio] = average
    return accuracyDist

dist = calcAccuracyDist()

# Plot the Results
x = dist.keys()
y = dist.values()
plt.plot(np.unique(x), np.poly1d(np.polyfit(x, y, 4))(np.unique(x)), label="Best Fit Line")
plt.scatter(x,y)
plt.xlabel("Ratio")
plt.ylabel("% Accuracy")
plt.show()


text:
import os
import argparse
import csv
import cv2
import numpy as np
import pybullet
import uuid
import pdb
from PIL import Image
import time
import generate_urdf as urdf
import math

class ImageGenerator:
    image_size = (256, 256)  # [px]

    camera_distance = 18.0  # [m]
    pixel_size = 3.25  # [px/m]

    rotation_distance = 2 * camera_distance  # [m]
    object_types = ['box', 'cylinder', 'sphere']

    def __init__(self, data_directory, image_directory, rollout_len, segmented_directory, objects_directory='/tmp/generated'):
        self.data_dir = data_directory
        self.objects_directory = objects_directory
        self.rollout_len = rollout_len
        self.image_directory = image_directory
        self.segm_directory = segmented_directory

        # Start pybullet 'DIRECT'ly without GUI
        pybullet.connect(pybullet.GUI)
        pybullet.setGravity(0, 0, -1000.0)
        pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_SHADOWS, 0)

    def generate_random_object(self):
        # Create and save random object
        object_type = np.random.choice(self.object_types)
        if object_type == 'cylinder':
            urdf_string, params = urdf.generate_cylinder()
            height = params[1]
            params = np.multiply(params, 100)
            file_name = 'cylinder_r_{}_h_{}.urdf'.format(int(params[0]), int(params[1]))
        elif object_type == 'box':
            urdf_string, params = urdf.generate_box()
            height = params[2]
            params = np.multiply(params, 100)
            file_name = 'box_x_{}_y_{}_z_{}.urdf'.format(int(params[0]), int(params[1]), int(params[2]))
        elif object_type == 'sphere':
            urdf_string, params = urdf.generate_sphere()
            height = 2 * params[0]
            params = np.multiply(params, 100)
            file_name = 'sphere_r_{}.urdf'.format(int(params[0]))

        return urdf_string, file_name, height

    def generate_simulated_images_for_random_scene(self):
        #pybullet.loadURDF('objects/plane.urdf', [0, 0, 0], useFixedBase=True)
        num_objects = np.random.choice([1, 2, 3])
        for i in range(num_objects):
            collision = False
            urdf_string, file_name, height = self.generate_random_object()
            file_path = os.path.join(self.objects_directory, file_name)

            with open(file_path, 'w') as file:
                file.write(urdf_string)

            position = [0.0, 0.0, 0.0]
            #position[0] = np.random.uniform(-6,6)
            #position[1] = np.random.uniform(-6,6)
            position[2] = height / 2  # insert the object at half object height so it does not collide with the plane

            orientation = [0.0, 0.0, 0.0]
            #orientation[0] = np.random.uniform(0, 360)
            #orientation[1] = np.random.uniform(0, 360)
            # orientation[2] = np.random.uniform(0, 360)

            if i > 0:
                distance = math.sqrt(np.sum((np.array(position)-old_position)**2))
                if distance < (height+old_height)/2:
                    collision = True
            old_position = np.array(position)
            old_height = height
            if not collision:
                model_id = pybullet.loadURDF(file_path, position, pybullet.getQuaternionFromEuler(orientation))
        # camera settings

        camera_position = np.random.uniform(0.0, 1.0, size=(3))
        camera_orientation = [0.0, 0.0, 0.0]  # pitch, roll, yaw [rad]
        camera_field_of_view = 60

        near_plane = 1.0
        far_plane = 100.0
        orientation = [0.0, 0.0, 0.0]
        orientation[0] = np.random.uniform(0, 360)
        orientation[1] = np.random.uniform(0, 360)

        #pybullet.resetBasePositionAndOrientation(model_id,  np.random.uniform(0.0, 1.0, size=(3)), pybullet.getQuaternionFromEuler(orientation))
        for i in range(1):
            pybullet.stepSimulation()

        depth_images = []
        segm_images = []
        first = True
        yaw = camera_orientation[2]
        pitch = -np.pi / 2
        roll = camera_orientation[1]
        pose_deltas = []
        start_time = time.time()
        dt = 0.05 #measured
        COEFF = 1.5
        frame = np.random.choice([7.0, 8.0, 11])
        print("Frame size: {}".format(frame))
        for i in range(0, self.rollout_len):
            delta = [0.0, 0.0, 0.0]
            pos, orn = pybullet.getBasePositionAndOrientation(model_id)
            if i == 0:
                obj_delta =[np.random.choice([0.5, 0.0, -0.5]), np.random.choice([0.5, 0.0, -0.5])]
            if not frame > pos[0] > -frame or not frame > pos[1] > -frame:
                obj_delta[0]*=-1
                obj_delta[1]*=-1

                newpos = [pos[0] + obj_delta[0], pos[1] + obj_delta[1], pos[2]]
            else:
                newpos = [pos[0]+ obj_delta[0], pos[1]+ obj_delta[1], pos[2]]
            #print(newpos)
            pybullet.resetBasePositionAndOrientation(model_id, newpos, orn)
            #print(start_time-time.time())
            start_time=time.time()
            if first:
                delta = [0.0, 0.0, 0.0]
                delta_pos = [0.0, 0.0, 0.0]
                first = True
            else:
                delta[0] = np.clip(pose_deltas[-1][3] + math.sqrt(dt)* np.random.randn(), -0.00001, 0.00001)
                delta[1] = np.clip(pose_deltas[-1][4] + math.sqrt(dt)* np.random.randn(), -0.001, 0.001)
                delta[2] = np.clip(pose_deltas[-1][5] + math.sqrt(dt)* np.random.randn(), -0.00, 0.00)
                delta_pos[0] = np.clip(COEFF*delta_pos[0] + math.sqrt(dt)* np.random.randn(), -0.1, 0.1)
                delta_pos[1] = np.clip(COEFF*delta_pos[1] + math.sqrt(dt)* np.random.randn(), -0.1, 0.1)
                delta_pos[2] = np.clip(COEFF*delta_pos[2] + math.sqrt(dt)* np.random.randn(), -0.02, 0.02)
            yaw += delta[2]
            pitch +=delta[0]
            roll += delta[1]
            camera_position_candidate =camera_position + delta_pos
            if not 3.0 > camera_position_candidate[0] > -3.0:
                delta_pos[0] = 0.0
            if not 3.0 > camera_position_candidate[1] > -3.0:
                delta_pos[1] = 0.0
            if not 3.0 > camera_position_candidate[2] > -3.0:
                delta_pos[2] = 0.0

            camera_position += delta_pos
            view_matrix = pybullet.computeViewMatrixFromYawPitchRoll(
                camera_position,
                self.camera_distance,
                np.rad2deg(yaw),
                np.rad2deg(pitch),
                np.rad2deg(roll),
                2
            )
            projection_matrix = pybullet.computeProjectionMatrixFOV(
                camera_field_of_view,
                self.image_size[0] / self.image_size[1],
                near_plane,
                far_plane
            )

            image = pybullet.getCameraImage(
                self.image_size[0],
                self.image_size[1],
                view_matrix,
                projection_matrix,
                shadow=0,
                lightDirection=[1, 1, 1],
                renderer=pybullet.ER_BULLET_HARDWARE_OPENGL
            )
            segm_image = image[4]
            color_image = np.reshape(np.array(image[2])[:, :, :3], (image[1], image[0], 3))

            segm_images.append(segm_image)
            depth_images.append(color_image)
            full_delta = np.concatenate((delta_pos, delta), axis=0)
            #print(full_delta)
            #print(camera_position)
            pose_deltas.append(full_delta)
        pybullet.resetSimulation()
        return depth_images, pose_deltas, segm_images

    def generate_images_for_random_scene(self, num, combine_images=True):
        images, poses, segm_images = self.generate_simulated_images_for_random_scene()
        id = uuid.uuid4()

        # Front image
        j=0

        imgs_folder = os.path.join(self.image_directory, "rollout_"+str(num))
        segm_imgs_folder = os.path.join(self.segm_directory, "rollout_"+str(num))

        if not os.path.isdir(imgs_folder):
            os.mkdir(imgs_folder)
        if not os.path.isdir(segm_imgs_folder):
            os.mkdir(segm_imgs_folder)
        for image, pose, segm_image in zip(images, poses, segm_images):

            depth_image_path = os.path.join(imgs_folder,'sim{}_rgb_{}.jpg'.format(num, j))
            img = Image.fromarray(image.astype(np.uint8))

            segm_image_path = os.path.join(segm_imgs_folder,'sim{}_sgm_{}.jpg'.format(num, j))
            segm_img = Image.fromarray(segm_image.astype(np.uint8))
            img.save(depth_image_path)
            segm_img.save(segm_image_path)
            j+=1
            yield [j, pose, depth_image_path, segm_image_path]

    def generate_images(self, number_simulations):
        for i in range(0, number_simulations):
            csv_name = "database"+str(i)+".csv"
            database_file = os.path.join(self.data_dir, "csv_files", csv_name)
            csv_file = open(database_file, mode='a')
            writer = csv.writer(csv_file)

            # If file is empty, write a header row
            if csv_file.tell() == 0:
                writer.writerow(['id', 'pose', 'image_path', 'segmented_image_path'])

            for row in gen.generate_images_for_random_scene(i):  # First value should be 0.0
                writer.writerow(row)

            csv_file.close()

if __name__ == "__main__":
    a = argparse.ArgumentParser()
    a.add_argument("--dir", help="Output path")
    a.add_argument("--len_roll", help="Length of rollout", default=60)
    a.add_argument("--num_roll", help="Number of rollouts", default=1000)

    args = a.parse_args()
    out_dir = args.dir
    img_dir = os.path.join(out_dir, "imgs64")
    segm_imgs = os.path.join(out_dir, "segm_imgs")
    csv_dir = os.path.join(out_dir, "csv_files")
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    if not os.path.exists(img_dir):
        os.mkdir(img_dir)
    if not os.path.exists(segm_imgs):
        os.mkdir(segm_imgs)
    if not os.path.exists(csv_dir):
        os.mkdir(csv_dir)

    len_rollout = int(args.len_roll)
    num_rollouts = int(args.num_roll)
    gen = ImageGenerator(out_dir, img_dir, len_rollout, segm_imgs)
    gen.generate_images(number_simulations=num_rollouts)