import os
import numpy as np
import pandas as pd
import h5py
import logging


logger = logging.getLogger(__name__)

# MAX_EDGE_READS = 500000000
# MAX_EDGE_READS = 200000000


def _get_names(index_type):
    if index_type.lower() in ['target', 'target_id', 'target_node_id', 'target_node_ids']:
        col_to_index = 'target_node_id'
        index_grp_name = 'indices/target_to_source'
    elif index_type in ['source', 'source_id', 'source_node_id', 'source_node_ids']:
        col_to_index = 'source_node_id'
        index_grp_name = 'indices/source_to_target'
    elif index_type in ['edge_type', 'edge_type_id', 'edge_type_ids']:
        col_to_index = 'edge_type_id'
        index_grp_name = 'indices/edge_type_to_index'
    else:
        raise ValueError('Unknown edges parameter {}'.format(index_type))

    return col_to_index, index_grp_name


def remove_index(edges_file, edges_population):
    with h5py.File(edges_file, mode='r+') as edges_h5:
        edges_pop_grp = edges_h5[edges_population]
        del edges_pop_grp['indices']


def create_index_in_memory(edges_file, edges_population, index_type, force_rebuild=True, compression='gzip', **kwargs):
    col_to_index, index_grp_name = _get_names(index_type)

    with h5py.File(edges_file, mode='r+') as edges_h5:
        edges_pop_grp = edges_h5[edges_population]

        if index_grp_name in edges_pop_grp:
            # Remove existing index if it exists
            if not force_rebuild:
                logger.debug('create_index_in_memory> Edges index {} already exists, skipping.'.format(index_grp_name))
            else:
                logger.debug('create_index_in_memory> Removing existing index {}.'.format(index_grp_name))
                del edges_pop_grp[index_grp_name]

        index_grp = edges_pop_grp.create_group(index_grp_name)
        ids_array = np.array(edges_pop_grp[col_to_index])  # ids to be indexed
        total_edges = len(edges_pop_grp[col_to_index])

        if total_edges == 0:
            logger.warning('edges file {} does not contain any edges.'.format(edges_file))
            return

        # group together and save ids so contigous duplicates are represented using ranges, creating a range -> edge
        # index table.
        #  eg. [10 10 10 10 10 32 32 32 10 10 ...] ==> ... 10: [(0, 5), (8, 10)], 32: [(5, 8)], ...
        logger.debug('create_index_in_memory> Creating range_to_edge_id table')
        ids_diffs = np.diff(ids_array)
        ids_diffs_idx = ids_diffs.nonzero()[0]

        ranges_beg = np.concatenate(([0], ids_diffs_idx + 1))
        ranges_end = np.concatenate((ids_diffs_idx + 1, [total_edges]))
        id_idxs = np.concatenate(([0], ids_diffs_idx + 1))

        r2e_table_df = pd.DataFrame({
            'lu_ids': ids_array[id_idxs],
            'range_beg': ranges_beg,
            'range_end': ranges_end
        }).sort_values('lu_ids')

        index_grp.create_dataset('range_to_edge_id', data=r2e_table_df[['range_beg', 'range_end']].values,
                                 dtype='uint64', compression=compression) # np.uint64)

        # create a map to the range_to_edge_id dataset from id --> blocks ranges. The id value is implicitly equal to
        # the index
        # TODO: See if del r2e_table_df will significantly improve memory footprint?
        logger.debug('create_index_in_memory> Creating node_id_to_range table')
        ordered_ids = np.array(r2e_table_df['lu_ids'])
        ordered_ids_diffs = np.diff(ordered_ids).nonzero()[0]
        ordered_ids_beg = np.concatenate(([0], ordered_ids_diffs+1))
        ordered_ids_end = np.concatenate((ordered_ids_diffs+1, [len(r2e_table_df)]))
        id_idxs = np.concatenate(([0], ordered_ids_diffs+1))
        i2r_table_df = pd.DataFrame({
            'lu_ids': ordered_ids[id_idxs],
            'range_beg': ordered_ids_beg,
            'range_end': ordered_ids_end
        }).set_index('lu_ids')

        # There may be missing values in the id_cols table so fil them in
        i2r_table_df = i2r_table_df.reindex(pd.RangeIndex(i2r_table_df.index.max()+1), fill_value=0)

        index_grp.create_dataset('node_id_to_range', data=i2r_table_df[['range_beg', 'range_end']].values,
                                 dtype=np.uint64, compression=compression)


def create_index_on_disk(edges_file, edges_population, index_type, force_rebuild=False, cache_file=None,
                         max_edge_reads=200000000, **kwargs):
    col_to_index, index_grp_name = _get_names(index_type)
    cache_file = cache_file if cache_file is not None else edges_file

    mode = 'r+' if os.path.exists(cache_file) else 'w'
    with h5py.File(cache_file, mode=mode) as cache_h5:
        # Step 1: Separate the column being indexed into N partitons each of less than max_edge_reads length. Build the
        # index for each partition and save to disk. Make sure we don't have to have the read the entire col_to_index
        # in memory at one time

        edges_h5 = cache_h5 if edges_file == cache_file else h5py.File(edges_file, 'r')
        edges_root_grp = edges_h5[edges_population]

        if edges_population not in cache_h5:
            caches_root_grp = cache_h5.create_group(edges_population)
        else:
            caches_root_grp = cache_h5[edges_population]

        total_edges = edges_root_grp[col_to_index].shape[0]
        partition_size = np.min((max_edge_reads, total_edges))
        n_partitions = np.ceil(total_edges / max_edge_reads).astype(np.uint)
        block_begin_idx = 0  # initial index of current partition being created
        partition_index = 0
        max_id = 0
        total_blocks = 0

        if force_rebuild:
            del caches_root_grp[index_grp_name]

        if index_grp_name not in caches_root_grp or 'cache' not in caches_root_grp[index_grp_name]:
            index_grp = caches_root_grp.create_group(index_grp_name) \
                if index_grp_name not in caches_root_grp else caches_root_grp[index_grp_name]
            cache_grp = index_grp.create_group('cache')
            cache_grp.attrs['max_edge_reads'] = max_edge_reads
            cache_grp.attrs['cache_partition_size'] = partition_size

        else:
            cache_grp = caches_root_grp[index_grp_name]['cache']
            max_id = cache_grp.attrs.get('max_id', max_id)

            if cache_grp.attrs['max_edge_reads'] != max_edge_reads:
                raise Exception(
                    'Cache already exists but has a conflicting max_edge_reads' 
                    '({} vs {}). Use force-rebuild'.format(cache_grp.attrs['max_edge_reads'], max_edge_reads)
                )

            if cache_grp.attrs['cache_partition_size'] != partition_size:
                raise Exception(
                    'Cache already exists but has a conflicting partition_size' 
                    '({} vs {}). Use force-rebuild'.format(cache_grp.attrs['cache_partition_size'], partition_size)
                )

        logger.debug('Number of edges, {}, exceeds maximum ({}).'.format(total_edges, max_edge_reads))
        logger.debug('Separating into {} partitions'.format(n_partitions))
        while block_begin_idx < total_edges:
            # cache_grp_name = index_grp_name + '/cache/edges_partition_{}'.format(partition_index)
            partition_grp_name = 'edges_partition_{}'.format(partition_index)
            if partition_grp_name in cache_grp:
                logger.debug('Cache {}/cache/{} already exists, skipping'.format(index_grp_name, partition_grp_name))
                block_begin_idx += partition_size
                partition_index += 1
                total_blocks += cache_grp[partition_grp_name]['partitioned_table'].shape[0]
                continue

            logger.debug('Creating cache {} of {} to {}/{}'.format(partition_index+1, n_partitions, index_grp_name,
                                                                   partition_grp_name))

            block_end_idx = block_begin_idx + partition_size
            edge_ids = edges_root_grp[col_to_index][block_begin_idx:block_end_idx]

            # Group together contigious duplicates to create a list of [id [edge_index_beg, edge_index_end)], sort by
            # id and save table to disk
            diffs = np.diff(edge_ids)
            diffs_idx = diffs.nonzero()[0]
            ranges_beg = np.concatenate(([0], diffs_idx + 1)) + block_begin_idx
            ranges_end = np.concatenate((diffs_idx + 1, [partition_size])) + block_begin_idx
            ranges_end[-1] = np.min((total_edges+1, ranges_end[-1]))
            id_idxs = np.concatenate(([0], diffs_idx + 1))

            r2e_table_df = pd.DataFrame({
                'lu_ids': edge_ids[id_idxs],
                'range_beg': ranges_beg,
                'range_end': ranges_end
            }).sort_values('lu_ids')

            # Creates a lookup table id --> ranges
            ordered_ids = r2e_table_df['lu_ids'].values
            # id_block_indxs = np.diff(ordered_ids).nonzero()[0] + 1
            # id_block_sizes = np.concatenate(([id_block_indxs[0]], np.diff(id_block_indxs)))
            # sums_vals = np.ones(len(ordered_ids), dtype=np.int32)
            # sums_vals[id_block_indxs] = -id_block_sizes + 1

            ordered_ids_diffs = np.diff(ordered_ids).nonzero()[0]
            ordered_ids_beg = np.concatenate(([0], ordered_ids_diffs + 1))
            ordered_ids_end = np.concatenate((ordered_ids_diffs + 1, [len(r2e_table_df)]))
            id_idxs = np.concatenate(([0], ordered_ids_diffs + 1))
            i2r_table_df = pd.DataFrame({
                'lu_ids': ordered_ids[id_idxs],
                'range_beg': ordered_ids_beg,
                'range_end': ordered_ids_end
            }).set_index('lu_ids')

            # fill in missing ids and foward fill Nans with the last previous index index
            i2r_table_df = i2r_table_df.reindex(pd.RangeIndex(i2r_table_df.index.max() + 1))
            i2r_table_df['range_end'] = i2r_table_df['range_end'].fillna(method='ffill')
            i2r_table_df['range_end'].fillna(0, inplace=True)

            nans_mask = i2r_table_df['range_beg'].isna()
            i2r_table_df['range_beg'][nans_mask] = i2r_table_df['range_end'][nans_mask]

            # Save partition to disk
            partition_grp = cache_grp.create_group(partition_grp_name)
            partition_grp.create_dataset('partitioned_table',
                                         data=r2e_table_df[['lu_ids', 'range_beg', 'range_end']].values,
                                         dtype='uint64')
            partition_grp.create_dataset('lookup_tables', data=i2r_table_df[['range_beg', 'range_end']].values,
                                         dtype='uint64')
            partition_grp.attrs['columns'] = 'id,range_beg,range_end'
            cache_h5.flush()

            max_id = np.max((max_id, int(np.max(ordered_ids))))
            total_blocks += len(r2e_table_df)
            block_begin_idx += partition_size
            partition_index += 1

            del r2e_table_df
            del i2r_table_df

        # TODO: Find max_id, total_blocks, and n_partitions and save as attribute in "cache" group
        cache_grp.attrs['max_id'] = max_id
        cache_grp.attrs['n_partitions'] = n_partitions
        cache_grp.attrs['total_blocks'] = total_blocks

    with h5py.File(cache_file, mode='r+') as cache_h5:
        # Step 2: Using the partitions we cached to disk in the previous step to create the actual index
        edges_pop_grp = cache_h5[edges_population]
        index_grp = edges_pop_grp[index_grp_name]
        if 'cache' not in index_grp:
            raise ValueError('Error, could not find hdf5 group {}/cache. Unable to proceed.'.format(index_grp_name))
        cache_grp = edges_pop_grp[index_grp_name]['cache']

        logger.debug('Building edges index {}'.format(index_grp_name))

        total_blocks = cache_grp.attrs['total_blocks']  # total number of 'partitioned_table' rows across all partitions
        max_id = cache_grp.attrs['max_id']  # maximum value in 'col_to_index' across all partions
        n_partitions = len(cache_grp.keys())

        # Go through each partition and build a master id lookup table. eg for each id should contains the number
        # of ranges and where their offsets would go when ordered.
        block_sizes = np.zeros(max_id + 1, dtype=np.uint32)
        for grp_name, grp in edges_pop_grp[index_grp_name]['cache'].items():
            # block_sizes += np.diff(grp['lookup_tables']).flatten()
            partiton_max_id = grp['lookup_tables'].shape[0]  # some partions may not contain all the ids
            block_sizes[:partiton_max_id] += np.diff(grp['lookup_tables']).flatten()

        # global_offsets = np.concatenate(([0], np.cumsum(block_sizes)[:-1])).astype(np.uint)  # + 1
        global_offsets = np.concatenate(([0], np.cumsum(block_sizes))).astype(np.uint)  # + 1
        ranges = np.vstack((global_offsets, np.concatenate((global_offsets[1:], [total_blocks])))).T

        # edges_pop_grp.create_dataset('{}/node_id_to_range'.format(index_grp), data=ranges, dtype=np.uint64)
        if not cache_grp.attrs.get('building_index', False):
            if 'node_id_to_range' in index_grp:
                del index_grp['node_id_to_range']
            index_grp.create_dataset('node_id_to_range'.format(index_grp), data=ranges[:-1], dtype=np.uint64)

            if 'range_to_edge_id' in index_grp:
                del index_grp['range_to_edge_id']
            index_grp.create_dataset('range_to_edge_id', (total_blocks, 2), dtype=np.uint64)

        # Using the ids offsets we can calculate the number N such that the associated "partitioned_table" for those N
        # ids have less than MAX_EDGE_READS rows (across all partitions). Procedure is to find N_0, iterate through all
        # the partitions to find ranges for ids [0, N_0], merge and sort and save to dataset. Then do the same thing
        # for ids [N_0, N_1]
        # TODO: Save index_beg, index_end, and block number as attributes in "cache"
        caches = {grp_name: grp for grp_name, grp in edges_pop_grp[index_grp_name]['cache'].items()}
        read_block_size = max_edge_reads / (n_partitions*2+1)

        if cache_grp.attrs.get('building_index', False):
            print('Index has already been partially built, resumming from last point')
            block_num = cache_grp.attrs['block_num']
            id_beg = cache_grp.attrs['id_beg']
            id_end = cache_grp.attrs['id_end']

        else:
            # separating global offsets into N blocks. Need to track which block were on, the begging index/row of the
            # block and the end index/row of block
            cache_grp.attrs['building_index'] = True
            cache_grp.attrs['block_num'] = block_num = 1
            cache_grp.attrs['id_beg'] = id_beg = 0
            id_end = np.max((np.searchsorted(global_offsets, read_block_size * block_num, side='right'), 1))
            id_end = np.min((max_id, id_end))
            cache_grp.attrs['id_end'] = id_end

        while id_beg <= max_id:
            logger.debug('Building "range_to_edge_id" for nodes ({}, {})'.format(id_beg, id_end))
            indx_beg = global_offsets[id_beg]
            indx_end = global_offsets[id_end+1]
            block_num += 1

            to_remove = []
            master_df = None
            for grp_name, grp in caches.items():
                # Use "lookup_tables" to find the which rows in "partioned_table" corresponds to ids [id_beg, id_end].
                # The partitioned tables are already sorted.
                if id_beg > grp['lookup_tables'].shape[0]:
                    to_remove.append(grp_name)
                    continue

                if id_end >= grp['lookup_tables'].shape[0]:
                    part_beg, part_end = grp['lookup_tables'][id_beg][0], grp['lookup_tables'][-1][1]
                    to_remove.append(grp_name)
                    # del caches[grp_name]
                else:
                    part_beg, part_end = grp['lookup_tables'][id_beg][0], grp['lookup_tables'][id_end][1]

                if part_end - part_beg <= 0:
                    continue  # possible partion does contain any edges for the subset of ids we are looking for

                tmp_df = pd.DataFrame({
                    'id': grp['partitioned_table'][part_beg:part_end, 0],
                    'range_beg': grp['partitioned_table'][part_beg:part_end, 1],
                    'range_end': grp['partitioned_table'][part_beg:part_end, 2]
                })
                master_df = tmp_df if master_df is None else pd.concat((master_df, tmp_df))

            if master_df is not None:
                master_df = master_df.sort_values(['id', 'range_beg'])
                index_grp['range_to_edge_id'][indx_beg:indx_end, :] = master_df[['range_beg', 'range_end']]

                for name in to_remove:
                    del caches[name]

                del master_df

            id_beg = id_end + 1
            id_end = np.searchsorted(global_offsets, read_block_size*block_num, side='right')
            id_end = np.min((max_id, id_end))

            # Save state of loop in case of failure and we need to progress later
            cache_grp.attrs['block_num'] = block_num
            cache_grp.attrs['id_beg'] = indx_beg
            cache_grp.attrs['id_end'] = indx_end

        # Remove cache
        if 'cache' in index_grp:
            del index_grp['cache']
