import collections as col
import sqlite3 

def get_2d_tileset_info(db_file):
    conn = sqlite3.connect(db_file)
    c = conn.cursor()

    row = c.execute("SELECT * from tileset_info").fetchone()
    tileset_info = {
        'zoom_step': row[0],
        'max_length': row[1],
        'assembly': row[2],
        'chrom_names': row[3],
        'chrom_sizes': row[4],
        'tile_size': row[5],
        'max_zoom': row[6],
        'max_width': row[7],
        'min_pos': [1, 1],
        'max_pos': [row[1], row[1]]
    }
    conn.close()

    return tileset_info

def get_1D_tiles(db_file, tile_ids):
    '''
    Retrieve a contiguous set of tiles from a 2D db tile file.

    Parameters
    ----------
    db_file: str
        The filename of the sqlite db file
    zoom: int
        The zoom level
    tile_x_pos: int
        The x position of the first tile
    numx: int
        The width of the block of tiles to retrieve

    Returns
    -------
    tiles: {pos: tile_value}
        A set of tiles, indexed by position
    '''
    tileset_info = get_2d_tileset_info(db_file)

    conn = sqlite3.connect(db_file)

    c = conn.cursor()
    tile_width = tileset_info['max_width'] / 2 ** zoom

    tile_x_start_pos = tile_width * tile_x_pos
    tile_x_end_pos = tile_x_start_pos + (numx * tile_width)

    # print('tile_x_start:', tile_x_start_pos, tile_x_end_pos)

    query = '''
    SELECT fromX, toX, fromY, toY, chrOffset, importance, fields, uid
    FROM intervals,position_index
    WHERE
        intervals.id=position_index.id AND
        zoomLevel <= {} AND
        rToX >= {} AND
        rFromX <= {}
    '''.format(
        zoom,
        tile_x_start_pos,
        tile_x_end_pos,
    )

    rows = c.execute(query).fetchall()

    new_rows = col.defaultdict(list)
    # print("len(rows)", len(rows))

    for r in rows:
        try:
            uid = r[7].decode('utf-8')
        except AttributeError:
            uid = r[7]

        x_start = r[0]
        x_end = r[1]
        y_start = r[2]
        y_end = r[3]

        for i in range(tile_x_pos, tile_x_pos + numx):
            tile_x_start = i * tile_width
            tile_x_end = (i+1) * tile_width

            if (
                x_start < tile_x_end and
                x_end >= tile_x_start
            ):
                # add the position offset to the returned values
                new_rows[i] += [
                    {'xStart': r[0],
                     'xEnd': r[1],
                     'yStart': r[2],
                     'yEnd': r[3],
                     'chrOffset': r[4],
                     'importance': r[5],
                     'uid': uid,
                     'fields': r[6].split('\t')}]
    conn.close()

    return new_rows

def get_2D_tiles(db_file, zoom, tile_x_pos, tile_y_pos, numx=1, numy=1):
    '''
    Retrieve a contiguous set of tiles from a 2D db tile file.

    Parameters
    ----------
    db_file: str
        The filename of the sqlite db file
    zoom: int
        The zoom level
    tile_x_pos: int
        The x position of the first tile
    tile_y_pos: int
        The y position of the first tile
    numx: int
        The width of the block of tiles to retrieve
    numy: int
        The height of the block of tiles to retrieve

    Returns
    -------
    tiles: {pos: tile_value}
        A set of tiles, indexed by position
    '''
    tileset_info = get_2d_tileset_info(db_file)

    conn = sqlite3.connect(db_file)

    c = conn.cursor()
    tile_width = tileset_info['max_width'] / 2 ** zoom

    tile_x_start_pos = tile_width * tile_x_pos
    tile_x_end_pos = tile_x_start_pos + (numx * tile_width)

    tile_y_start_pos = tile_width * tile_y_pos
    tile_y_end_pos = tile_y_start_pos + (numy * tile_width)

    query = '''
    SELECT fromX, toX, fromY, toY, chrOffset, importance, fields, uid
    FROM intervals,position_index
    WHERE
        intervals.id=position_index.id AND
        zoomLevel <= {} AND
        rToX >= {} AND
        rFromX <= {} AND
        rToY >= {} AND
        rFromY <= {}
    '''.format(
        zoom,
        tile_x_start_pos,
        tile_x_end_pos,
        tile_y_start_pos,
        tile_y_end_pos
    )

    rows = c.execute(query).fetchall()

    new_rows = col.defaultdict(list)

    for r in rows:
        try:
            uid = r[7].decode('utf-8')
        except AttributeError:
            uid = r[7]

        x_start = r[0]
        x_end = r[1]
        y_start = r[2]
        y_end = r[3]

        for i in range(tile_x_pos, tile_x_pos + numx):
            for j in range(tile_y_pos, tile_y_pos + numy):
                tile_x_start = i * tile_width
                tile_x_end = (i+1) * tile_width

                tile_y_start = j * tile_width
                tile_y_end = (j+1) * tile_width

                if (
                    x_start < tile_x_end and
                    x_end >= tile_x_start and
                    y_start < tile_y_end and
                    y_end >= tile_y_start
                ):
                    # add the position offset to the returned values
                    new_rows[(i, j)] += [
                        {'xStart': r[0],
                         'xEnd': r[1],
                         'yStart': r[2],
                         'yEnd': r[3],
                         'chrOffset': r[4],
                         'importance': r[5],
                         'uid': uid,
                         'fields': r[6].split('\t')}]
    conn.close()

    return new_rows
