<commit_before>#!/usr/bin/env python

"""
This module is for creating or editing a MagIC contribution,
(or a piece of one).
You can build a contribution or an individual table from the ground up,
or you can read in one or more MagIC-format files.
You can also extract specific data from a table --
for instance, you can build a DIblock for plotting.
"""
import os
import re
import numpy as np
import pandas as pd
from pandas import DataFrame
# from pmagpy import pmag
from pmagpy import data_model3 as data_model
from pmagpy import controlled_vocabularies3 as cv


class Contribution(object):

    """
    A Contribution is a collection of MagicDataFrames,
    each of which corresponds to one MagIC table.
    The Contribution object also has methods for
    manipulating one or more tables in the contribution --
    for example, renaming a site.
    """

    def __init__(self, directory=".", read_tables='all',
                 custom_filenames=None, single_file=None,
                 dmodel=None, vocabulary=""):
        if isinstance(dmodel, data_model.DataModel):
            self.data_model = dmodel
            Contribution.dmodel = dmodel
        else:
            try:
                self.data_model = Contribution.dmodel
            except AttributeError:
                Contribution.dmodel = data_model.DataModel()
                self.data_model = Contribution.dmodel
        if isinstance(vocabulary, cv.Vocabulary):
            self.vocab = vocabulary
        else:
            try:
                self.vocab = Contribution.vocab
            except AttributeError:
                Contribution.vocab = cv.Vocabulary(dmodel=self.data_model)
                self.vocab = Contribution.vocab
        self.directory = os.path.realpath(directory)
        self.table_names = ['measurements', 'specimens', 'samples',
                            'sites', 'locations', 'contribution',
                            'criteria', 'ages', 'images']
        self.ancestry = ['measurements', 'specimens', 'samples',
                         'sites', 'locations']

        # create standard filenames from table names
        self.filenames = {t: t + ".txt" for t in self.table_names}
        # update self.filenames to include custom names
        if custom_filenames:
            self.add_custom_filenames(custom_filenames)
        self.tables = {}
        # if not otherwise specified, read in all possible tables
        if read_tables == 'all':
            read_tables = self.table_names
        if single_file:  # use if filename is known but type isn't
            self.add_magic_table('unknown', single_file)
            return
        else:  # read in data for all required tables
            for name in read_tables:
                self.add_magic_table(name)

    ## Methods for building up the contribution

    def add_custom_filenames(self, custom_filenames):
        """
        Update/overwrite self.filenames with custom names.
        Input dict should have the format:
        {"specimens": "custom_spec_file.txt", ...}
        """
        if custom_filenames:
            self.filenames.update(custom_filenames)

    def add_empty_magic_table(self, dtype, col_names=None, groups=None):
        """
        Add a blank MagicDataFrame to the contribution.
        You can provide either a list of column names,
        or a list of column group names.
        If provided, col_names takes precedence.
        """
        if dtype not in self.table_names:
            print "-W- {} is not a valid MagIC table name".format(dtype)
            print "-I- Valid table names are: {}".format(", ".join(self.table_names))
            return
        data_container = MagicDataFrame(dtype=dtype, columns=col_names, groups=groups)
        self.tables[dtype] = data_container


    def add_magic_table(self, dtype, fname=None):
        """
        Read in a new file to add a table to self.tables.
        """
        # if providing a filename but no data type
        if dtype == "unknown":
            filename = os.path.join(self.directory, fname)
            data_container = MagicDataFrame(filename, dmodel=self.data_model)
            dtype = data_container.dtype
            if dtype == 'empty':
                return False
            else:
                self.tables[dtype] = data_container
                return data_container
        # if providing a data type, use the canonical filename
        elif dtype not in self.filenames:
            print '-W- "{}" is not a valid MagIC table type'.format(dtype)
            print "-I- Available table types are: {}".format(", ".join(self.table_names))
            return False
        filename = os.path.join(self.directory, self.filenames[dtype])
        if os.path.exists(filename):
            data_container = MagicDataFrame(filename, dmodel=self.data_model)
            if data_container.dtype != "empty":
                self.tables[dtype] = data_container
                return data_container
        else:
            print "-W- No such file: {}".format(filename)
            return False


    def propagate_measurement_info(self):
        """
        Take a contribution with a measurement table.
        Create specimen, sample, site, and location tables
        using the unique names in the measurement table to fill in
        the index.
        """
        meas_df = self.tables['measurements'].df
        names_list = ['specimen', 'sample', 'site', 'location']
        # add in any tables that you can
        for num, name in enumerate(names_list):
            # don't replace tables that already exist
            if (name + "s") in self.tables:
                continue
            elif name in meas_df.columns:
                print "making new {} file".format(name)
                items = meas_df[name].unique()
                df = pd.DataFrame(columns=[name], index=items)
                df[name] = df.index
                # add in parent name if possible
                # (i.e., sample name to specimens table)
                if num < (len(names_list) - 1):
                    parent = names_list[num+1]
                    if parent in meas_df.columns:
                        df[parent] = meas_df.drop_duplicates(subset=[name])[parent].values
                self.tables[name + "s"] = MagicDataFrame(dtype=name + "s", df=df)
                self.write_table_to_file(name + "s")


    def propagate_all_tables_info(self, write=True):
        """
        Find any items (specimens, samples, sites, or locations) from
        tables other than measurements and make sure they each have a
        row in their own table.  For example, if a site name is in
        the samples table but not in the sites table, create a row
        for it in the sites table.
        """
        for table_name in ["specimens", "samples", "sites", "locations"]:
            if not table_name in self.tables:
                continue
            df = self.tables[table_name].df
            parent_name, child_name = self.get_parent_and_child(table_name)
            if parent_name:
                if parent_name[:-1] in df.columns:
                    parents = sorted(set(df[parent_name[:-1]].dropna().values))
                    if parent_name in self.tables: # if there is a parent table, update it
                        parent_df = self.tables[parent_name].df
                        missing_parents = set(parents) - set(parent_df.index)
                        if missing_parents: # add any missing values
                            print "-I- Updating {} table with values from {} table".format(parent_name, table_name)
                            for item in missing_parents:
                                self.add_item(parent_name, {parent_name[:-1]: item}, label=item)
                            # save any changes to file
                            if write:
                                self.write_table_to_file(parent_name)

                    else:  # if there is no parent table, create it if necessary
                        if parents:
                            # create a parent_df with the names you got from the child
                            print "-I- Creating new {} table with data from {} table".format(parent_name, table_name)
                            parent_df = pd.DataFrame(columns=[parent_name[:-1]], index=parents)
                            parent_df[parent_name[:-1]] = parent_df.index
                            self.tables[parent_name] = MagicDataFrame(dtype=parent_name,
                                                                      df=parent_df)
                            if write:
                                # save new table to file
                                self.write_table_to_file(parent_name)
            if child_name:
                if child_name in df.columns:
                    raw_children = df[child_name].dropna().str.split(':')
                    # create dict of all children with parent info
                    parent_of_child = {}
                    for parent, children in raw_children.iteritems():
                        for child in children:
                            # remove whitespace
                            child = child.strip()
                            old_parent = parent_of_child.get(child)
                            if old_parent and (old_parent != parent):
                                print '-I- for {} {}, replacing: {} with: {}'.format(child_name, child, old_parent, parent)
                            parent_of_child[child] = parent
                    # old way:
                    # flatten list, ignore duplicates
                    #children = sorted(set([item.strip() for sublist in raw_children for item in sublist]))
                    if child_name in self.tables: # if there is already a child table, update it
                        child_df = self.tables[child_name].df
                        missing_children = set(parent_of_child.keys()) - set(child_df.index)
                        if missing_children: # add any missing values
                            print "-I- Updating {} table with values from {} table".format(child_name, table_name)
                            for item in missing_children:
                                data = {child_name[:-1]: item, table_name[:-1]: parent_of_child[item]}
                                self.add_item(child_name, data, label=item)
                            if write:
                                # save any changes to file
                                self.write_table_to_file(child_name)
                    else: # if there is no child table, create it if necessary
                        if children:
                            # create a child_df with the names you got from the parent
                            print "-I- Creating new {} table with data from {} table".format(child_name, table_name)
                            # old way to make new table:
                            #child_df = pd.DataFrame(columns=[table_name[:-1]], index=children)
                            # new way to make new table
                            children_list = sorted(parent_of_child.keys())
                            children_data = [[child_name, parent_of_child[c_name]] for c_name in children_list]
                            child_df = pd.DataFrame(index=children_list, columns=[child_name[:-1], table_name[:-1]], data=children_data)

                            self.tables[child_name] = MagicDataFrame(dtype=child_name, df=child_df)
                            if write:
                                # save new table to file
                                self.write_table_to_file(child_name)


    def get_parent_and_child(self, table_name):
        """
        Get the name of the parent table and the child table
        for a given MagIC table name.

        Parameters
        ----------
        table_name : string of MagIC table name ['specimens', 'samples', 'sites', 'locations']

        Returns
        -------
        parent_name : string of parent table name
        child_name : string of child table name
        """
        if table_name not in self.ancestry:
            return None, None
        parent_ind = self.ancestry.index(table_name) + 1
        if parent_ind + 1 > len(self.ancestry):
            parent_name = None
        else:
            parent_name = self.ancestry[parent_ind]
        child_ind = self.ancestry.index(table_name) - 1
        if child_ind < 0:
            child_name = None
        else:
            child_name = self.ancestry[child_ind]
        return parent_name, child_name

    def get_min_max_lat_lon(self):
        """
        Find latitude/longitude information from sites table
        and group it by location.

        Returns
        ---------
        """
        if 'sites' not in self.tables:
            return
        # get min/max lat/lon from sites table
        site_container = self.tables['sites']
        if not ('lat' in site_container.df.columns and 'lon' in site_container.df.columns):
            return
        # convert lat/lon columns to string type
        # (this is necessary for consistency because they MAY be string type already)
        site_container.df['lat'] = site_container.df['lat'].fillna('').astype(str)
        site_container.df['lon'] = site_container.df['lon'].fillna('').astype(str)
        # replace empty strings with np.nan
        site_container.df['lat'] = np.where(site_container.df['lat'].str.len(), site_container.df['lat'], np.nan)
        site_container.df['lon'] = np.where(site_container.df['lon'].str.len(), site_container.df['lon'], np.nan)
        # convert lat/lon values to float (they make be string from grid)
        site_container.df['lat'] = site_container.df['lat'].astype(float)
        site_container.df['lon'] = site_container.df['lon'].astype(float)
        # group lat/lon by location
        grouped_lon = site_container.df[['lon', 'location']].groupby('location')
        grouped_lat = site_container.df[['lat', 'location']].groupby('location')
        # get min/max longitude by location
        lon_w = grouped_lon.min()
        lon_e = grouped_lon.max()
        # get min/max latitude by location
        lat_s = grouped_lat.min()
        lat_n = grouped_lat.max()
        # assign lat/lon to location table
        locs = {}
        if 'locations' not in self.tables:
            return
        loc_container = self.tables['locations']
        for loc in lat_s.index:
            coords = {}
            coords['lat_s'] = lat_s.loc[loc]['lat']
            coords['lat_n'] = lat_n.loc[loc]['lat']
            coords['lon_e'] = lon_e.loc[loc]['lon']
            coords['lon_w'] = lon_w.loc[loc]['lon']
            locs[loc] = coords
        loc_container = self.tables['locations']
        for loc_name in locs:
            if loc_name in loc_container.df.index:
                coords = locs[loc_name]
                for coord in locs[loc_name]:
                    # warn user if an old value will be overwritten
                    new_value = coords[coord]
                    if coord not in loc_container.df.columns:
                        loc_container.df[coord] = None
                    old_value = loc_container.df.ix[loc_name, coord]
                    # use first value if multiple values returned, but don't shorten a string
                    if not (isinstance(old_value, str) or isinstance(old_value, unicode)):
                        try:
                            old_value = old_value[0]
                        except TypeError: # if only one value
                            pass
                        except IndexError: # if np.nan
                            pass
                    if old_value is None:
                        pass
                    elif isinstance(old_value, str) or isinstance(old_value, unicode):
                        try:
                            old_value = float(old_value)
                        except ValueError:
                            print '-W- In {}, automatically generated {} value ({}) will overwrite previous value ({})'.format(loc_name, coord, new_value, old_value)
                            old_value = None
                    elif np.isnan(old_value):
                        pass
                    elif new_value != old_value:
                        print '-W- In {}, automatically generated {} value ({}) will overwrite previous value ({})'.format(loc_name, coord, new_value, old_value)
                    # set new value
                    loc_container.df.set_value(loc_name, coord, new_value)
        self.write_table_to_file('locations')
        return locs

    def add_item(self, table_name, data, label):
        self.tables[table_name].add_row(label, data)

    ## Methods for making changes to a Contribution
    ## that need to propagate to multiple tables

    def rename_item(self, table_name, item_old_name, item_new_name):
        """
        Rename item (such as a site) everywhere that it occurs.
        This change often spans multiple tables.
        For example, a site name will occur in the sites table,
        the samples table, and possibly in the locations/ages tables.
        """
        # define some helper methods:

        def put_together_if_list(item):
            """
            String joining function
            that doesn't break with None/np.nan
            """
            try:
                res = ":".join(item)
                return ":".join(item)
            except TypeError as ex:
                #print ex
                return item

        def replace_colon_delimited_value(df, col_name, old_value, new_value):
            """
            Col must contain list
            """
            count = 1
            for index, row in df[df[col_name].notnull()].iterrows():
                names_list = row[col_name]
                names_list = [name.strip() for name in names_list]
                try:
                    ind = names_list.index(old_value)
                except ValueError as ex:
                    count += 1
                    continue
                names_list[ind] = new_value
                df.ix[count, col_name] = names_list
                count += 1

        # initialize some things
        item_type = table_name
        ###col_name = item_type[:-1] + "_name"
        col_name = item_type[:-1]
        col_name_plural = col_name + "s"
        table_df = self.tables[item_type].df

        if item_old_name == '':
            # just add a new item
            self.add_item(table_name, {col_name: item_new_name}, item_new_name)
            return

        # rename item in its own table
        table_df.rename(index={item_old_name: item_new_name}, inplace=True)
        # rename in any parent/child tables
        for table_name in self.tables:
            df = self.tables[table_name].df
            col_names = df.columns
            # change anywhere col_name (singular, i.e. site) is found
            if col_name in col_names:
                df[col_name].where(df[col_name] != item_old_name, item_new_name, inplace=True)
                # change anywhere col_name (plural, i.e. sites) is found
            if col_name_plural in col_names:
                df[col_name_plural + "_list"] = df[col_name_plural].str.split(":")
                replace_colon_delimited_value(df, col_name_plural + "_list", item_old_name, item_new_name)
                df[col_name_plural] = df[col_name_plural + "_list"].apply(put_together_if_list)
                df.drop(col_name_plural + "_list", axis=1, inplace=True)
            self.tables[table_name].df = df


    def get_table_name(self, ind):
        """
        Return both the table_name (i.e., 'specimens')
        ###and the col_name (i.e., 'specimen_name')
        and the col_name (i.e., 'specimen')
        for a given index in self.ancestry.
        """
        if ind >= len(self.ancestry):
            return "", ""
        if ind > -1:
            table_name = self.ancestry[ind]
            ###name = table_name[:-1] + "_name"
            name = table_name[:-1]
            return table_name, name
        return "", ""


    def propagate_name_down(self, col_name, df_name):
        """
        Put the data for "col_name" into dataframe with df_name
        Used to add 'site_name' to specimen table, for example.
        """
        if df_name not in self.tables:
            self.add_magic_table(df_name)
        df = self.tables[df_name].df
        if col_name in df.columns:
            print '{} already in {}'.format(col_name, df_name)
            return df

        # otherwise, do necessary merges to get col_name into df
        # get names for each level
        grandparent_table_name = col_name.split('_')[0] + "s"
        grandparent_name = grandparent_table_name[:-1]
        ind = self.ancestry.index(grandparent_table_name) - 1
        #
        parent_table_name, parent_name = self.get_table_name(ind)
        child_table_name, child_name = self.get_table_name(ind - 1)
        bottom_table_name, bottom_name = self.get_table_name(ind - 2)

        # merge in bottom level
        if child_name not in df.columns:
            # add child table if missing
            if bottom_table_name not in self.tables:
                result = self.add_magic_table(bottom_table_name)
                if not isinstance(result, MagicDataFrame):
                    print "-W- Couldn't read in {} data".format(bottom_table_name)
                    print "-I- Make sure you've provided the correct file name"
                    return df
            # add child_name to df
            add_df = self.tables[bottom_table_name].df
            # drop duplicate names
            add_df = add_df.drop_duplicates(subset=bottom_name)
            df = df.merge(add_df[[child_name]],
                          left_on=[bottom_name],
                          right_index=True, how="left")
            self.tables[df_name].df = df

        # merge in one level above
        if parent_name not in df.columns:
            # add parent_table if missing
            if child_table_name not in self.tables:
                result = self.add_magic_table(child_table_name)
                if not isinstance(result, MagicDataFrame):
                    print "-W- Couldn't read in {} data".format(child_table_name)
                    print "-I- Make sure you've provided the correct file name"
                    return df
            # add parent_name to df
            add_df = self.tables[child_table_name].df
            # drop duplicate names
            add_df = add_df.drop_duplicates(subset=child_name)
            df = df.merge(add_df[[parent_name]],
                          left_on=[child_name],
                          right_index=True, how="left")
            self.tables[df_name].df = df

        # merge in two levels above
        if grandparent_name not in df.columns:
            # add grandparent table if it is missing
            if parent_table_name not in self.tables:
                result = self.add_magic_table(parent_table_name)
                if not isinstance(result, MagicDataFrame):
                    print "-W- Couldn't read in {} data".format(parent_table_name)
                    print "-I- Make sure you've provided the correct file name"
                    return df
            # add grandparent name to df
            add_df = self.tables[parent_table_name].df
            # drop duplicate names
            add_df = add_df.drop_duplicates(subset=parent_name)
            df = df.merge(add_df[[grandparent_name]],
                          left_on=[parent_name],
                          right_index=True, how="left")
        # update the Contribution
        self.tables[df_name].df = df
        return df


    def propagate_cols_down(self, col_names, target_df_name, source_df_name):
        """
        Put the data for "col_name" from source_df into target_df
        Used to get "azimuth" from sample table into measurements table
        (for example).
        Note: if getting data from the sample table, don't include "sample"
        in the col_names list.  It is included automatically.
        """
        # make sure target table is read in
        if target_df_name not in self.tables:
            self.add_magic_table(target_df_name)
        if target_df_name not in self.tables:
            print "-W- Couldn't read in {} table".format(target_df_name)
            return
        # make sure source table is read in
        if source_df_name not in self.tables:
            self.add_magic_table(source_df_name)
            print "-W- Couldn't read in {} table".format(source_df_name)
            return
        # make sure col_names are all available in source table
        source_df = self.tables[source_df_name].df
        if not set(col_names).issubset(source_df.columns):
            for col in col_names[:]:
                if col not in source_df.columns:
                    print "-W- Column '{}' isn't in {} table, skipping it".format(col, source_df_name)
                    col_names.remove(col)
        if not col_names:
            print "-W- Invalid or missing column names, could not propagate down"
            return

        add_name = source_df_name[:-1]
        self.propagate_name_down(add_name, target_df_name)
        #
        target_df = self.tables[target_df_name].df
        source_df = self.tables[source_df_name].df
        #
        target_df = target_df.merge(source_df[col_names], how='left',
                                    left_on=add_name, right_index=True,
                                    suffixes=["_target", "_source"])
        # mess with target_df to remove unneded merge columns
        for col in col_names:
            # if there has been a previous merge, consolidate and delete data
            if col + "_target" in target_df.columns:
                # prioritize values from target df
                new_arr = np.where(target_df[col + "_target"],
                                   target_df[col + "_target"],
                                   target_df[col + "_source"])
                target_df.rename(columns={col + "_target": col}, inplace=True)
                target_df[col] = new_arr
            if col + "_source" in target_df.columns:
                # delete extra merge column
                del target_df[col + "_source"]
        #
        self.tables[target_df_name].df = target_df
        return target_df

    def remove_non_magic_cols(self):
        """
        Remove all non-MagIC columns from all tables.
        """
        for table_name in self.tables:
            table = self.tables[table_name]
            table_dm = self.data_model.dm[table_name]
            approved_cols = table_dm.index
            unrecognized_cols = (set(table.df.columns) - set(approved_cols))
            if unrecognized_cols:
                print '-I- Removing non-MagIC column names from {}:'.format(table_name),
                for col in unrecognized_cols:
                    self.tables[table_name].df.drop(col, axis='columns', inplace=True)
                    print col,
                print "\n"

    def write_table_to_file(self, dtype):
        """
        Write out a MagIC table to file, using custom filename
        as specified in self.filenames.

        Parameters
        ----------
        dtype : str
            magic table name
        """
        fname = self.filenames[dtype]
        if dtype in self.tables:
            self.tables[dtype].write_magic_file(custom_name=fname,
                                                dir_path=self.directory)



class MagicDataFrame(object):

    """
    Each MagicDataFrame corresponds to one MagIC table.
    The MagicDataFrame object consists of a pandas DataFrame,
    and assorted methods for manipulating that DataFrame.
    """

    def __init__(self, magic_file=None, columns=None, dtype=None,
                 groups=None, dmodel=None, df=None):
        """
        Provide either a magic_file or a dtype.
        List of columns is optional,
        and will only be used if magic_file == None.
        Instead of a list of columns, you can also provide
        a list of group-names, and the specific col_names
        will be filled in by the data model.
        If provided, col_names takes precedence.
        """
        if isinstance(df, pd.DataFrame):
            self.df = df
            if dtype:
                self.dtype = dtype
            else:
                print '-W- Please provide data type...'
        # make sure all required arguments are present
        if not magic_file and not dtype and not isinstance(df, pd.DataFrame):
            print "-W- To make a MagicDataFrame, you must provide either a filename or a datatype"
            self.df = None
            return
        # fetch data model if not provided
        # (DataModel type sometimes not recognized, hence ugly hack below)
        if isinstance(dmodel, data_model.DataModel) or str(data_model.DataModel) == str(type(dmodel)):
            MagicDataFrame.data_model = dmodel
            self.data_model = dmodel
        else:
            try:
                self.data_model = MagicDataFrame.data_model
            except AttributeError:
                MagicDataFrame.data_model = data_model.DataModel()
                self.data_model = MagicDataFrame.data_model

        if isinstance(df, pd.DataFrame):
            pass
        # if no file is provided, make an empty dataframe of the appropriate type
        elif not magic_file:
            self.dtype = dtype
            if not isinstance(columns, type(None)):
                self.df = DataFrame(columns=columns)
            else:
                self.df = DataFrame()
                self.df.index.name = dtype[:-1] if dtype.endswith("s") else dtype
        # if there is a file provided, read in the data and ascertain dtype
        else:
            ## new way of reading in data using pd.read_table
            with open(magic_file) as f:
                try:
                    delim, dtype = f.readline().split('\t')[:2]
                except ValueError as ex:
                    print ex, type(ex)
                    print "-W- Empty file {}".format(magic_file)
                    self.dtype = 'empty'
                    self.df = DataFrame()
                    return
            self.df = pd.read_table(magic_file, skiprows=[0])
            self.dtype = dtype.strip()
            if self.dtype.endswith('s'):
                name = '{}'.format(self.dtype[:-1])
            else:
                name = self.dtype
            if self.dtype == 'measurements':
                self.df['measurement'] = self.df['experiment'] + self.df['number'].astype(str)
            elif self.dtype == 'contribution':
                name = 'doi'
                return
            elif self.dtype == 'images':
                return
            elif self.dtype == 'criteria':
                self.df.index = self.df['table_column']
                return
            #
            if len(self.df) and self.dtype != 'ages':
                self.df.index = self.df[name].astype(str)
            elif self.dtype == 'ages':
                self.df.index = self.df.index.astype(str)
            #del self.df[name]
            #self.dtype = dtype
            # replace '' with None, so you can use isnull(), notnull(), etc.
            # can always switch back with DataFrame.fillna('')
            self.df = self.df.where(self.df.notnull(), None)

            # drop any completely blank columns
            # this is not necessarily a good idea....
            #self.df.dropna(axis=1, how='all', inplace=True)
            #
            # add df columns that were passed in but weren't in the file
            if columns:
                for col in columns:
                    if col not in self.df.columns:
                        self.df[col] = None

        # add col_names by group
        if groups and not columns:
            columns = []
            for group_name in groups:
                columns.extend(list(self.data_model.get_group_headers(self.dtype, group_name)))
            for col in columns:
                if col not in self.df.columns:
                    self.df[col] = None
            self.df = self.df[columns]
        # make sure type column is present (i.e., sample column in samples df)
        self.df[self.dtype[:-1]] = self.df.index


    ## Methods to change self.df inplace

    def update_row(self, ind, row_data):
        """
        Update a row with data.
        Must provide the specific numeric index (not row label).
        If any new keys are present in row_data dictionary,
        that column will be added to the dataframe.
        This is done inplace.
        """
        if sorted(row_data.keys()) != sorted(self.df.columns):<commit_msg>fix bug with age column vs index in ages table<commit_after>#!/usr/bin/env python

"""
This module is for creating or editing a MagIC contribution,
(or a piece of one).
You can build a contribution or an individual table from the ground up,
or you can read in one or more MagIC-format files.
You can also extract specific data from a table --
for instance, you can build a DIblock for plotting.
"""
import os
import re
import numpy as np
import pandas as pd
from pandas import DataFrame
# from pmagpy import pmag
from pmagpy import data_model3 as data_model
from pmagpy import controlled_vocabularies3 as cv


class Contribution(object):

    """
    A Contribution is a collection of MagicDataFrames,
    each of which corresponds to one MagIC table.
    The Contribution object also has methods for
    manipulating one or more tables in the contribution --
    for example, renaming a site.
    """

    def __init__(self, directory=".", read_tables='all',
                 custom_filenames=None, single_file=None,
                 dmodel=None, vocabulary=""):
        if isinstance(dmodel, data_model.DataModel):
            self.data_model = dmodel
            Contribution.dmodel = dmodel
        else:
            try:
                self.data_model = Contribution.dmodel
            except AttributeError:
                Contribution.dmodel = data_model.DataModel()
                self.data_model = Contribution.dmodel
        if isinstance(vocabulary, cv.Vocabulary):
            self.vocab = vocabulary
        else:
            try:
                self.vocab = Contribution.vocab
            except AttributeError:
                Contribution.vocab = cv.Vocabulary(dmodel=self.data_model)
                self.vocab = Contribution.vocab
        self.directory = os.path.realpath(directory)
        self.table_names = ['measurements', 'specimens', 'samples',
                            'sites', 'locations', 'contribution',
                            'criteria', 'ages', 'images']
        self.ancestry = ['measurements', 'specimens', 'samples',
                         'sites', 'locations']

        # create standard filenames from table names
        self.filenames = {t: t + ".txt" for t in self.table_names}
        # update self.filenames to include custom names
        if custom_filenames:
            self.add_custom_filenames(custom_filenames)
        self.tables = {}
        # if not otherwise specified, read in all possible tables
        if read_tables == 'all':
            read_tables = self.table_names
        if single_file:  # use if filename is known but type isn't
            self.add_magic_table('unknown', single_file)
            return
        else:  # read in data for all required tables
            for name in read_tables:
                self.add_magic_table(name)

    ## Methods for building up the contribution

    def add_custom_filenames(self, custom_filenames):
        """
        Update/overwrite self.filenames with custom names.
        Input dict should have the format:
        {"specimens": "custom_spec_file.txt", ...}
        """
        if custom_filenames:
            self.filenames.update(custom_filenames)

    def add_empty_magic_table(self, dtype, col_names=None, groups=None):
        """
        Add a blank MagicDataFrame to the contribution.
        You can provide either a list of column names,
        or a list of column group names.
        If provided, col_names takes precedence.
        """
        if dtype not in self.table_names:
            print "-W- {} is not a valid MagIC table name".format(dtype)
            print "-I- Valid table names are: {}".format(", ".join(self.table_names))
            return
        data_container = MagicDataFrame(dtype=dtype, columns=col_names, groups=groups)
        self.tables[dtype] = data_container


    def add_magic_table(self, dtype, fname=None):
        """
        Read in a new file to add a table to self.tables.
        """
        # if providing a filename but no data type
        if dtype == "unknown":
            filename = os.path.join(self.directory, fname)
            data_container = MagicDataFrame(filename, dmodel=self.data_model)
            dtype = data_container.dtype
            if dtype == 'empty':
                return False
            else:
                self.tables[dtype] = data_container
                return data_container
        # if providing a data type, use the canonical filename
        elif dtype not in self.filenames:
            print '-W- "{}" is not a valid MagIC table type'.format(dtype)
            print "-I- Available table types are: {}".format(", ".join(self.table_names))
            return False
        filename = os.path.join(self.directory, self.filenames[dtype])
        if os.path.exists(filename):
            data_container = MagicDataFrame(filename, dmodel=self.data_model)
            if data_container.dtype != "empty":
                self.tables[dtype] = data_container
                return data_container
        else:
            print "-W- No such file: {}".format(filename)
            return False


    def propagate_measurement_info(self):
        """
        Take a contribution with a measurement table.
        Create specimen, sample, site, and location tables
        using the unique names in the measurement table to fill in
        the index.
        """
        meas_df = self.tables['measurements'].df
        names_list = ['specimen', 'sample', 'site', 'location']
        # add in any tables that you can
        for num, name in enumerate(names_list):
            # don't replace tables that already exist
            if (name + "s") in self.tables:
                continue
            elif name in meas_df.columns:
                print "making new {} file".format(name)
                items = meas_df[name].unique()
                df = pd.DataFrame(columns=[name], index=items)
                df[name] = df.index
                # add in parent name if possible
                # (i.e., sample name to specimens table)
                if num < (len(names_list) - 1):
                    parent = names_list[num+1]
                    if parent in meas_df.columns:
                        df[parent] = meas_df.drop_duplicates(subset=[name])[parent].values
                self.tables[name + "s"] = MagicDataFrame(dtype=name + "s", df=df)
                self.write_table_to_file(name + "s")


    def propagate_all_tables_info(self, write=True):
        """
        Find any items (specimens, samples, sites, or locations) from
        tables other than measurements and make sure they each have a
        row in their own table.  For example, if a site name is in
        the samples table but not in the sites table, create a row
        for it in the sites table.
        """
        for table_name in ["specimens", "samples", "sites", "locations"]:
            if not table_name in self.tables:
                continue
            df = self.tables[table_name].df
            parent_name, child_name = self.get_parent_and_child(table_name)
            if parent_name:
                if parent_name[:-1] in df.columns:
                    parents = sorted(set(df[parent_name[:-1]].dropna().values))
                    if parent_name in self.tables: # if there is a parent table, update it
                        parent_df = self.tables[parent_name].df
                        missing_parents = set(parents) - set(parent_df.index)
                        if missing_parents: # add any missing values
                            print "-I- Updating {} table with values from {} table".format(parent_name, table_name)
                            for item in missing_parents:
                                self.add_item(parent_name, {parent_name[:-1]: item}, label=item)
                            # save any changes to file
                            if write:
                                self.write_table_to_file(parent_name)

                    else:  # if there is no parent table, create it if necessary
                        if parents:
                            # create a parent_df with the names you got from the child
                            print "-I- Creating new {} table with data from {} table".format(parent_name, table_name)
                            parent_df = pd.DataFrame(columns=[parent_name[:-1]], index=parents)
                            parent_df[parent_name[:-1]] = parent_df.index
                            self.tables[parent_name] = MagicDataFrame(dtype=parent_name,
                                                                      df=parent_df)
                            if write:
                                # save new table to file
                                self.write_table_to_file(parent_name)
            if child_name:
                if child_name in df.columns:
                    raw_children = df[child_name].dropna().str.split(':')
                    # create dict of all children with parent info
                    parent_of_child = {}
                    for parent, children in raw_children.iteritems():
                        for child in children:
                            # remove whitespace
                            child = child.strip()
                            old_parent = parent_of_child.get(child)
                            if old_parent and (old_parent != parent):
                                print '-I- for {} {}, replacing: {} with: {}'.format(child_name, child, old_parent, parent)
                            parent_of_child[child] = parent
                    # old way:
                    # flatten list, ignore duplicates
                    #children = sorted(set([item.strip() for sublist in raw_children for item in sublist]))
                    if child_name in self.tables: # if there is already a child table, update it
                        child_df = self.tables[child_name].df
                        missing_children = set(parent_of_child.keys()) - set(child_df.index)
                        if missing_children: # add any missing values
                            print "-I- Updating {} table with values from {} table".format(child_name, table_name)
                            for item in missing_children:
                                data = {child_name[:-1]: item, table_name[:-1]: parent_of_child[item]}
                                self.add_item(child_name, data, label=item)
                            if write:
                                # save any changes to file
                                self.write_table_to_file(child_name)
                    else: # if there is no child table, create it if necessary
                        if children:
                            # create a child_df with the names you got from the parent
                            print "-I- Creating new {} table with data from {} table".format(child_name, table_name)
                            # old way to make new table:
                            #child_df = pd.DataFrame(columns=[table_name[:-1]], index=children)
                            # new way to make new table
                            children_list = sorted(parent_of_child.keys())
                            children_data = [[child_name, parent_of_child[c_name]] for c_name in children_list]
                            child_df = pd.DataFrame(index=children_list, columns=[child_name[:-1], table_name[:-1]], data=children_data)

                            self.tables[child_name] = MagicDataFrame(dtype=child_name, df=child_df)
                            if write:
                                # save new table to file
                                self.write_table_to_file(child_name)


    def get_parent_and_child(self, table_name):
        """
        Get the name of the parent table and the child table
        for a given MagIC table name.

        Parameters
        ----------
        table_name : string of MagIC table name ['specimens', 'samples', 'sites', 'locations']

        Returns
        -------
        parent_name : string of parent table name
        child_name : string of child table name
        """
        if table_name not in self.ancestry:
            return None, None
        parent_ind = self.ancestry.index(table_name) + 1
        if parent_ind + 1 > len(self.ancestry):
            parent_name = None
        else:
            parent_name = self.ancestry[parent_ind]
        child_ind = self.ancestry.index(table_name) - 1
        if child_ind < 0:
            child_name = None
        else:
            child_name = self.ancestry[child_ind]
        return parent_name, child_name

    def get_min_max_lat_lon(self):
        """
        Find latitude/longitude information from sites table
        and group it by location.

        Returns
        ---------
        """
        if 'sites' not in self.tables:
            return
        # get min/max lat/lon from sites table
        site_container = self.tables['sites']
        if not ('lat' in site_container.df.columns and 'lon' in site_container.df.columns):
            return
        # convert lat/lon columns to string type
        # (this is necessary for consistency because they MAY be string type already)
        site_container.df['lat'] = site_container.df['lat'].fillna('').astype(str)
        site_container.df['lon'] = site_container.df['lon'].fillna('').astype(str)
        # replace empty strings with np.nan
        site_container.df['lat'] = np.where(site_container.df['lat'].str.len(), site_container.df['lat'], np.nan)
        site_container.df['lon'] = np.where(site_container.df['lon'].str.len(), site_container.df['lon'], np.nan)
        # convert lat/lon values to float (they make be string from grid)
        site_container.df['lat'] = site_container.df['lat'].astype(float)
        site_container.df['lon'] = site_container.df['lon'].astype(float)
        # group lat/lon by location
        grouped_lon = site_container.df[['lon', 'location']].groupby('location')
        grouped_lat = site_container.df[['lat', 'location']].groupby('location')
        # get min/max longitude by location
        lon_w = grouped_lon.min()
        lon_e = grouped_lon.max()
        # get min/max latitude by location
        lat_s = grouped_lat.min()
        lat_n = grouped_lat.max()
        # assign lat/lon to location table
        locs = {}
        if 'locations' not in self.tables:
            return
        loc_container = self.tables['locations']
        for loc in lat_s.index:
            coords = {}
            coords['lat_s'] = lat_s.loc[loc]['lat']
            coords['lat_n'] = lat_n.loc[loc]['lat']
            coords['lon_e'] = lon_e.loc[loc]['lon']
            coords['lon_w'] = lon_w.loc[loc]['lon']
            locs[loc] = coords
        loc_container = self.tables['locations']
        for loc_name in locs:
            if loc_name in loc_container.df.index:
                coords = locs[loc_name]
                for coord in locs[loc_name]:
                    # warn user if an old value will be overwritten
                    new_value = coords[coord]
                    if coord not in loc_container.df.columns:
                        loc_container.df[coord] = None
                    old_value = loc_container.df.ix[loc_name, coord]
                    # use first value if multiple values returned, but don't shorten a string
                    if not (isinstance(old_value, str) or isinstance(old_value, unicode)):
                        try:
                            old_value = old_value[0]
                        except TypeError: # if only one value
                            pass
                        except IndexError: # if np.nan
                            pass
                    if old_value is None:
                        pass
                    elif isinstance(old_value, str) or isinstance(old_value, unicode):
                        try:
                            old_value = float(old_value)
                        except ValueError:
                            print '-W- In {}, automatically generated {} value ({}) will overwrite previous value ({})'.format(loc_name, coord, new_value, old_value)
                            old_value = None
                    elif np.isnan(old_value):
                        pass
                    elif new_value != old_value:
                        print '-W- In {}, automatically generated {} value ({}) will overwrite previous value ({})'.format(loc_name, coord, new_value, old_value)
                    # set new value
                    loc_container.df.set_value(loc_name, coord, new_value)
        self.write_table_to_file('locations')
        return locs

    def add_item(self, table_name, data, label):
        self.tables[table_name].add_row(label, data)

    ## Methods for making changes to a Contribution
    ## that need to propagate to multiple tables

    def rename_item(self, table_name, item_old_name, item_new_name):
        """
        Rename item (such as a site) everywhere that it occurs.
        This change often spans multiple tables.
        For example, a site name will occur in the sites table,
        the samples table, and possibly in the locations/ages tables.
        """
        # define some helper methods:

        def put_together_if_list(item):
            """
            String joining function
            that doesn't break with None/np.nan
            """
            try:
                res = ":".join(item)
                return ":".join(item)
            except TypeError as ex:
                #print ex
                return item

        def replace_colon_delimited_value(df, col_name, old_value, new_value):
            """
            Col must contain list
            """
            count = 1
            for index, row in df[df[col_name].notnull()].iterrows():
                names_list = row[col_name]
                names_list = [name.strip() for name in names_list]
                try:
                    ind = names_list.index(old_value)
                except ValueError as ex:
                    count += 1
                    continue
                names_list[ind] = new_value
                df.ix[count, col_name] = names_list
                count += 1

        # initialize some things
        item_type = table_name
        ###col_name = item_type[:-1] + "_name"
        col_name = item_type[:-1]
        col_name_plural = col_name + "s"
        table_df = self.tables[item_type].df

        if item_old_name == '':
            # just add a new item
            self.add_item(table_name, {col_name: item_new_name}, item_new_name)
            return

        # rename item in its own table
        table_df.rename(index={item_old_name: item_new_name}, inplace=True)
        # rename in any parent/child tables
        for table_name in self.tables:
            df = self.tables[table_name].df
            col_names = df.columns
            # change anywhere col_name (singular, i.e. site) is found
            if col_name in col_names:
                df[col_name].where(df[col_name] != item_old_name, item_new_name, inplace=True)
                # change anywhere col_name (plural, i.e. sites) is found
            if col_name_plural in col_names:
                df[col_name_plural + "_list"] = df[col_name_plural].str.split(":")
                replace_colon_delimited_value(df, col_name_plural + "_list", item_old_name, item_new_name)
                df[col_name_plural] = df[col_name_plural + "_list"].apply(put_together_if_list)
                df.drop(col_name_plural + "_list", axis=1, inplace=True)
            self.tables[table_name].df = df


    def get_table_name(self, ind):
        """
        Return both the table_name (i.e., 'specimens')
        ###and the col_name (i.e., 'specimen_name')
        and the col_name (i.e., 'specimen')
        for a given index in self.ancestry.
        """
        if ind >= len(self.ancestry):
            return "", ""
        if ind > -1:
            table_name = self.ancestry[ind]
            ###name = table_name[:-1] + "_name"
            name = table_name[:-1]
            return table_name, name
        return "", ""


    def propagate_name_down(self, col_name, df_name):
        """
        Put the data for "col_name" into dataframe with df_name
        Used to add 'site_name' to specimen table, for example.
        """
        if df_name not in self.tables:
            self.add_magic_table(df_name)
        df = self.tables[df_name].df
        if col_name in df.columns:
            print '{} already in {}'.format(col_name, df_name)
            return df

        # otherwise, do necessary merges to get col_name into df
        # get names for each level
        grandparent_table_name = col_name.split('_')[0] + "s"
        grandparent_name = grandparent_table_name[:-1]
        ind = self.ancestry.index(grandparent_table_name) - 1
        #
        parent_table_name, parent_name = self.get_table_name(ind)
        child_table_name, child_name = self.get_table_name(ind - 1)
        bottom_table_name, bottom_name = self.get_table_name(ind - 2)

        # merge in bottom level
        if child_name not in df.columns:
            # add child table if missing
            if bottom_table_name not in self.tables:
                result = self.add_magic_table(bottom_table_name)
                if not isinstance(result, MagicDataFrame):
                    print "-W- Couldn't read in {} data".format(bottom_table_name)
                    print "-I- Make sure you've provided the correct file name"
                    return df
            # add child_name to df
            add_df = self.tables[bottom_table_name].df
            # drop duplicate names
            add_df = add_df.drop_duplicates(subset=bottom_name)
            df = df.merge(add_df[[child_name]],
                          left_on=[bottom_name],
                          right_index=True, how="left")
            self.tables[df_name].df = df

        # merge in one level above
        if parent_name not in df.columns:
            # add parent_table if missing
            if child_table_name not in self.tables:
                result = self.add_magic_table(child_table_name)
                if not isinstance(result, MagicDataFrame):
                    print "-W- Couldn't read in {} data".format(child_table_name)
                    print "-I- Make sure you've provided the correct file name"
                    return df
            # add parent_name to df
            add_df = self.tables[child_table_name].df
            # drop duplicate names
            add_df = add_df.drop_duplicates(subset=child_name)
            df = df.merge(add_df[[parent_name]],
                          left_on=[child_name],
                          right_index=True, how="left")
            self.tables[df_name].df = df

        # merge in two levels above
        if grandparent_name not in df.columns:
            # add grandparent table if it is missing
            if parent_table_name not in self.tables:
                result = self.add_magic_table(parent_table_name)
                if not isinstance(result, MagicDataFrame):
                    print "-W- Couldn't read in {} data".format(parent_table_name)
                    print "-I- Make sure you've provided the correct file name"
                    return df
            # add grandparent name to df
            add_df = self.tables[parent_table_name].df
            # drop duplicate names
            add_df = add_df.drop_duplicates(subset=parent_name)
            df = df.merge(add_df[[grandparent_name]],
                          left_on=[parent_name],
                          right_index=True, how="left")
        # update the Contribution
        self.tables[df_name].df = df
        return df


    def propagate_cols_down(self, col_names, target_df_name, source_df_name):
        """
        Put the data for "col_name" from source_df into target_df
        Used to get "azimuth" from sample table into measurements table
        (for example).
        Note: if getting data from the sample table, don't include "sample"
        in the col_names list.  It is included automatically.
        """
        # make sure target table is read in
        if target_df_name not in self.tables:
            self.add_magic_table(target_df_name)
        if target_df_name not in self.tables:
            print "-W- Couldn't read in {} table".format(target_df_name)
            return
        # make sure source table is read in
        if source_df_name not in self.tables:
            self.add_magic_table(source_df_name)
            print "-W- Couldn't read in {} table".format(source_df_name)
            return
        # make sure col_names are all available in source table
        source_df = self.tables[source_df_name].df
        if not set(col_names).issubset(source_df.columns):
            for col in col_names[:]:
                if col not in source_df.columns:
                    print "-W- Column '{}' isn't in {} table, skipping it".format(col, source_df_name)
                    col_names.remove(col)
        if not col_names:
            print "-W- Invalid or missing column names, could not propagate down"
            return

        add_name = source_df_name[:-1]
        self.propagate_name_down(add_name, target_df_name)
        #
        target_df = self.tables[target_df_name].df
        source_df = self.tables[source_df_name].df
        #
        target_df = target_df.merge(source_df[col_names], how='left',
                                    left_on=add_name, right_index=True,
                                    suffixes=["_target", "_source"])
        # mess with target_df to remove unneded merge columns
        for col in col_names:
            # if there has been a previous merge, consolidate and delete data
            if col + "_target" in target_df.columns:
                # prioritize values from target df
                new_arr = np.where(target_df[col + "_target"],
                                   target_df[col + "_target"],
                                   target_df[col + "_source"])
                target_df.rename(columns={col + "_target": col}, inplace=True)
                target_df[col] = new_arr
            if col + "_source" in target_df.columns:
                # delete extra merge column
                del target_df[col + "_source"]
        #
        self.tables[target_df_name].df = target_df
        return target_df

    def remove_non_magic_cols(self):
        """
        Remove all non-MagIC columns from all tables.
        """
        for table_name in self.tables:
            table = self.tables[table_name]
            table_dm = self.data_model.dm[table_name]
            approved_cols = table_dm.index
            unrecognized_cols = (set(table.df.columns) - set(approved_cols))
            if unrecognized_cols:
                print '-I- Removing non-MagIC column names from {}:'.format(table_name),
                for col in unrecognized_cols:
                    self.tables[table_name].df.drop(col, axis='columns', inplace=True)
                    print col,
                print "\n"

    def write_table_to_file(self, dtype):
        """
        Write out a MagIC table to file, using custom filename
        as specified in self.filenames.

        Parameters
        ----------
        dtype : str
            magic table name
        """
        fname = self.filenames[dtype]
        if dtype in self.tables:
            self.tables[dtype].write_magic_file(custom_name=fname,
                                                dir_path=self.directory)



class MagicDataFrame(object):

    """
    Each MagicDataFrame corresponds to one MagIC table.
    The MagicDataFrame object consists of a pandas DataFrame,
    and assorted methods for manipulating that DataFrame.
    """

    def __init__(self, magic_file=None, columns=None, dtype=None,
                 groups=None, dmodel=None, df=None):
        """
        Provide either a magic_file or a dtype.
        List of columns is optional,
        and will only be used if magic_file == None.
        Instead of a list of columns, you can also provide
        a list of group-names, and the specific col_names
        will be filled in by the data model.
        If provided, col_names takes precedence.
        """
        if isinstance(df, pd.DataFrame):
            self.df = df
            if dtype:
                self.dtype = dtype
            else:
                print '-W- Please provide data type...'
        # make sure all required arguments are present
        if not magic_file and not dtype and not isinstance(df, pd.DataFrame):
            print "-W- To make a MagicDataFrame, you must provide either a filename or a datatype"
            self.df = None
            return
        # fetch data model if not provided
        # (DataModel type sometimes not recognized, hence ugly hack below)
        if isinstance(dmodel, data_model.DataModel) or str(data_model.DataModel) == str(type(dmodel)):
            MagicDataFrame.data_model = dmodel
            self.data_model = dmodel
        else:
            try:
                self.data_model = MagicDataFrame.data_model
            except AttributeError:
                MagicDataFrame.data_model = data_model.DataModel()
                self.data_model = MagicDataFrame.data_model

        if isinstance(df, pd.DataFrame):
            pass
        # if no file is provided, make an empty dataframe of the appropriate type
        elif not magic_file:
            self.dtype = dtype
            if not isinstance(columns, type(None)):
                self.df = DataFrame(columns=columns)
            else:
                self.df = DataFrame()
                self.df.index.name = dtype[:-1] if dtype.endswith("s") else dtype
        # if there is a file provided, read in the data and ascertain dtype
        else:
            ## new way of reading in data using pd.read_table
            with open(magic_file) as f:
                try:
                    delim, dtype = f.readline().split('\t')[:2]
                except ValueError as ex:
                    print ex, type(ex)
                    print "-W- Empty file {}".format(magic_file)
                    self.dtype = 'empty'
                    self.df = DataFrame()
                    return
            self.df = pd.read_table(magic_file, skiprows=[0])
            self.dtype = dtype.strip()
            if self.dtype.endswith('s'):
                name = '{}'.format(self.dtype[:-1])
            else:
                name = self.dtype
            if self.dtype == 'measurements':
                self.df['measurement'] = self.df['experiment'] + self.df['number'].astype(str)
            elif self.dtype == 'contribution':
                name = 'doi'
                return
            elif self.dtype == 'images':
                return
            elif self.dtype == 'criteria':
                self.df.index = self.df['table_column']
                return
            #
            if len(self.df) and self.dtype != 'ages':
                self.df.index = self.df[name].astype(str)
            elif self.dtype == 'ages':
                self.df.index = self.df.index.astype(str)
            #del self.df[name]
            #self.dtype = dtype
            # replace '' with None, so you can use isnull(), notnull(), etc.
            # can always switch back with DataFrame.fillna('')
            self.df = self.df.where(self.df.notnull(), None)

            # drop any completely blank columns
            # this is not necessarily a good idea....
            #self.df.dropna(axis=1, how='all', inplace=True)
            #
            # add df columns that were passed in but weren't in the file
            if columns:
                for col in columns:
                    if col not in self.df.columns:
                        self.df[col] = None

        # add col_names by group
        if groups and not columns:
            columns = []
            for group_name in groups:
                columns.extend(list(self.data_model.get_group_headers(self.dtype, group_name)))
            for col in columns:
                if col not in self.df.columns:
                    self.df[col] = None
            self.df = self.df[columns]
        # make sure type column is present (i.e., sample column in samples df)
        if self.dtype != 'ages':
            self.df[self.dtype[:-1]] = self.df.index


    ## Methods to change self.df inplace

    def update_row(self, ind, row_data):
        """
        Update a row with data.
        Must provide the specific numeric index (not row label).
        If any new keys are present in row_data dictionary,
        that column will be added to the dataframe.
        This is done inplace.
        """
        if sorted(row_data.keys()) != sorted(self.df.columns):
{'fraction_non_alphanumeric': 0.08640437490505848, 'fraction_numerical': 0.0012152514051344372, 'mean_word_length': 1.6677067714876201, 'pattern_counts': {'":': 14, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3267935', 'n_tokens_mistral': 17704, 'n_tokens_neox': 17076, 'n_words': 5093}