from __future__ import division, print_function
import json

import numpy as np
import pandas as pd
import cooler
import h5py
import time

TILESIZE = 256

where = np.flatnonzero
chromsizes = cooler.read_chromsizes(
    "http://s3.amazonaws.com/pkerp/data/mm9/chromInfo.txt"
)  # defaults to reading chr#,X,Y,M
chromosomes = list(chromsizes.keys())
chromid_map = dict(zip(chromosomes, range(len(chromosomes))))
cumul_lengths = np.r_[0, np.cumsum(chromsizes)]


def absCoord2bin(c, pos):
    try:
        cid = where(cumul_lengths > pos)[0] - 1
    except IndexError:
        return c.info["nbins"]
    chrom = chromosomes[cid]
    relPos = pos - cumul_lengths[cid]
    return c.offset((chrom, relPos, chromsizes[chrom]))


def getData(FILEPATH, zoomLevel, startPos1, endPos1, startPos2, endPos2):

    groupname = str(zoomLevel)

    with h5py.File(FILEPATH, "r") as f:
        c = cooler.Cooler(f[groupname])

        i0 = absCoord2bin(c, startPos1)
        i1 = absCoord2bin(c, endPos1)
        j0 = absCoord2bin(c, startPos2)
        j1 = absCoord2bin(c, endPos2)
        mat = c.matrix(balance=True)[i0:i1, j0:j1]

    flat = list(mat.toarray().ravel())

    return json.dumps({"dense": flat})


def getData2(cooler_matrix, zoomLevel, startPos1, endPos1, startPos2, endPos2):

    c = cooler_matrix["cooler"]
    matrix = cooler_matrix["matrix"]

    i0 = absCoord2bin(c, startPos1)
    i1 = absCoord2bin(c, endPos1)
    j0 = absCoord2bin(c, startPos2)
    j1 = absCoord2bin(c, endPos2)

    if (i1 - i0) == 0 or (j1 - j0) == 0:
        return pd.DataFrame(columns=["genome_start", "genome_end", "balanced"])

    t1 = time.time()
    pixels = matrix[i0:i1, j0:j1]
    cid1 = pixels["chrom1"].apply(chromid_map.get)
    cid2 = pixels["chrom2"].apply(chromid_map.get)
    pixels["genome_start"] = cumul_lengths[cid1] + pixels["start1"]
    pixels["genome_end"] = cumul_lengths[cid2] + pixels["end2"]

    print("z:", zoomLevel, "t:", time.time() - t1)
    return pixels[["genome_start", "genome_end", "balanced"]]


def getData3(cooler_matrix, zoomLevel, startPos1, endPos1, startPos2, endPos2):
    c = cooler_matrix["cooler"]

    i0 = absCoord2bin(c, startPos1)
    i1 = absCoord2bin(c, endPos1)
    j0 = absCoord2bin(c, startPos2)
    j1 = absCoord2bin(c, endPos2)

    if (i1 - i0) == 0 or (j1 - j0) == 0:
        return pd.DataFrame(columns=["genome_start", "genome_end", "balanced"])

    pixels = c.matrix(as_pixels=True, max_chunk=np.inf)[i0:i1, j0:j1]

    if not len(pixels):
        return pd.DataFrame(columns=["genome_start", "genome_end", "balanced"])

    lo = min(i0, j0)
    hi = max(i1, j1)
    bins = c.bins()[["chrom", "start", "end", "weight"]][lo:hi]
    bins["chrom"] = bins["chrom"].cat.codes
    pixels = cooler.annotate(pixels, bins)
    pixels["genome_start"] = cumul_lengths[pixels["chrom1"]] + pixels["start1"]
    pixels["genome_end"] = cumul_lengths[pixels["chrom2"]] + pixels["end2"]
    pixels["balanced"] = pixels["count"] * pixels["weight1"] * pixels["weight2"]

    return pixels[["genome_start", "genome_end", "balanced"]]


def getInfo(FILEPATH):

    with h5py.File(FILEPATH, "r") as f:
        total_length = int(cumul_lengths[-1])
        binsize = int(f["0"].attrs["bin-size"])
        binsize = 1000
        n_tiles = total_length / binsize / TILESIZE
        print("total_length:", total_length, binsize, TILESIZE)
        n_zooms = int(np.ceil(np.log2(n_tiles)))
        max_width = binsize * TILESIZE * 2 ** n_zooms

        info = {
            "min_pos": [0.0, 0.0],
            "max_pos": [total_length, total_length],
            "max_zoom": n_zooms,
            "max_width": max_width,
            "bins_per_dimension": TILESIZE,
        }

    return info

    return json.dumps(info)