import os
import os.path as op
import tempfile
from io import BytesIO

import h5py
import numpy as np
import pandas as pd
import pytest
from _common import isolated_filesystem

import cooler
import cooler.create

tmp = tempfile.gettempdir()
testdir = op.dirname(op.realpath(__file__))
datadir = op.join(testdir, "data")


@pytest.mark.parametrize(
    "fp", [op.join(datadir, "hg19.GM12878-MboI.matrix.2000kb.cool")]
)
def test_create_append(fp):
    import dask.dataframe as dd

    c = cooler.Cooler(fp)
    # chromsizes = c.chromsizes
    bins = c.bins()[:]
    pixels = c.pixels()[:]

    # create
    cooler.create.create(op.join(tmp, "test.df.2000kb.cool"), bins, pixels)
    cooler.create.create(
        op.join(tmp, "test.dict.2000kb.cool"),
        bins,
        {k: v for k, v in pixels.items()},
    )
    cooler.create.create(op.join(tmp, "test.iter_df.2000kb.cool"), bins, [pixels])
    cooler.create.create(
        op.join(tmp, "test.iter_dict.2000kb.cool"),
        bins,
        [{k: v for k, v in pixels.items()}],
    )
    ddf = dd.from_pandas(pixels, npartitions=3)
    cooler.create.create(op.join(tmp, "test.ddf.2000kb.cool"), bins, ddf)

    # Append
    cooler.create.append(
        op.join(tmp, "test.df.2000kb.cool"),
        "bins",
        {"start_1based": bins.apply(lambda x: x.start + 1, axis=1)},
    )
    cooler.create.append(op.join(tmp, "test.df.2000kb.cool"), "bins", {"ones": 1})
    series = ddf["count"] / ddf["count"].sum()
    series.name = "normed"
    cooler.create.append(op.join(tmp, "test.df.2000kb.cool"), "pixels", series)
    cooler.create.append(
        op.join(tmp, "test.df.2000kb.cool"), "pixels", series, force=True
    )
    cooler.create.append(
        op.join(tmp, "test.df.2000kb.cool"),
        "bins",
        {"twos": [np.ones(1000, dtype=int) * 2, np.ones(561, dtype=int) * 2]},
        chunked=True,
        force=True,
    )
    c2 = cooler.Cooler(op.join(tmp, "test.df.2000kb.cool"))
    assert len(c2.bins().columns) == 6
    assert len(c2.pixels().columns) == 4


@pytest.mark.parametrize(
    "f_hm,f_cool",
    [
        (
            op.join(datadir, "hg19.IMR90-MboI.matrix.2000kb.npy"),
            op.join(tmp, "test.cool"),
        )
    ],
)
@pytest.mark.skip("This test needs Internet connectivity.")
def test_roundtrip(f_hm, f_cool):
    chromsizes = cooler.read_chromsizes(
        "http://genome.ucsc.edu/goldenpath/help/hg19.chrom.sizes",
        name_patterns=(r"^chr[0-9]+$", r"chrX$"),
    )
    binsize = 2000000
    bintable = cooler.binnify(chromsizes, binsize)

    heatmap = np.load(f_hm)
    reader = cooler.create.ArrayLoader(bintable, heatmap, 100000)
    cooler.create.create(f_cool, bintable, reader, assembly="hg19")

    h5 = h5py.File(f_cool, "r")
    new_chromtable = cooler.api.chroms(h5)
    assert np.all(chromsizes.index == new_chromtable["name"])

    new_bintable = cooler.api.bins(h5)
    assert np.all(bintable == new_bintable)

    info = cooler.api.info(h5)
    assert info["genome-assembly"] == "hg19"
    assert info["bin-type"] == "fixed"
    assert info["bin-size"] == binsize

    mat = cooler.api.matrix(h5, 0, 100, 0, 100, "count", balance=False)
    assert mat.shape == (100, 100)
    assert np.allclose(heatmap[:100, :100], mat)

    mat = cooler.Cooler(h5).matrix("count", balance=False)[:100, :100]
    assert mat.shape == (100, 100)
    assert np.allclose(heatmap[:100, :100], mat)

    mat = cooler.api.matrix(h5, 100, 200, 100, 200, "count", balance=False)
    assert mat.shape == (100, 100)
    assert np.allclose(heatmap[100:200, 100:200], mat)

    mat = cooler.Cooler(h5).matrix("count", balance=False)[100:200, 100:200]
    assert mat.shape == (100, 100)
    assert np.allclose(heatmap[100:200, 100:200], mat)

    try:
        os.remove(f_cool)
    except OSError:
        pass


def test_rename_chroms():
    from shutil import copyfile

    with isolated_filesystem():
        copyfile(op.join(datadir, "toy.asymm.4.cool"), "toy.asymm.4.cool")
        clr = cooler.Cooler("toy.asymm.4.cool")
        assert clr.chromnames == ["chr1", "chr2"]
        cooler.rename_chroms(clr, {"chr1": "1", "chr2": "2"})
        assert clr.chromnames == ["1", "2"]  # the Cooler object is refreshed


def test_create_custom_cols():
    with isolated_filesystem():
        df = pd.DataFrame(
            {
                "bin1_id": [0, 1, 1, 1, 2, 2, 3, 4, 5],
                "bin2_id": [1, 1, 3, 4, 5, 6, 7, 8, 9],
                "foo": [1, 1, 1, 1, 1, 2, 2, 2, 2],
                "bar": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
            },
            columns=["bin1_id", "bin2_id", "foo", "bar"],
        )
        bins = pd.DataFrame(
            {
                "chrom": ["chr1"] * 5 + ["chr2"] * 5,
                "start": list(range(5)) * 2,
                "end": list(range(1, 6)) * 2,
            }
        )
        # works in unordered mode
        cooler.create_cooler("test.cool", bins, df, columns=["foo", "bar"])
        clr = cooler.Cooler("test.cool")
        assert len(clr.pixels().columns) == 4
        assert np.allclose(df, clr.pixels()[["bin1_id", "bin2_id", "foo", "bar"]][:])

        # works in ordered mode
        cooler.create_cooler(
            "test.cool", bins, df, columns=["foo", "bar"], ordered=True
        )
        clr = cooler.Cooler("test.cool")
        assert len(clr.pixels().columns) == 4
        assert np.allclose(df, clr.pixels()[["bin1_id", "bin2_id", "foo", "bar"]][:])

        # raises if no custom columns specified and 'count' does not exist
        with pytest.raises(ValueError):
            cooler.create_cooler("test.cool", bins, df, columns=None, ordered=True)


def test_write():
    chroms = pd.DataFrame(
        {
            "name": ["chr1", "chr2", "chr3"],
            "length": [32, 48, 56],
        },
        columns=["name", "length"],
    )
    bins = pd.DataFrame(
        {
            "chrom": [
                "chr1",
                "chr1",
                "chr2",
                "chr2",
                "chr2",
                "chr3",
                "chr3",
                "chr3",
                "chr3",
            ],
            "start": [0, 16, 0, 16, 32, 0, 16, 32, 48],
            "end": [16, 32, 16, 32, 48, 16, 32, 48, 56],
        },
        columns=["chrom", "start", "end"],
    )
    pixels = pd.DataFrame(
        {
            "bin1_id": [0],
            "bin2_id": [5],
            "count": [1],
        },
        columns=["bin1_id", "bin2_id", "count"],
    )
    h5opts = {"compression": "gzip", "compression_opts": 6}

    # chroms
    b = BytesIO()
    f = h5py.File(b, "r+")
    grp = f.create_group("chroms")
    cooler.create._create.write_chroms(grp, chroms, h5opts)
    f.flush()
    assert "name" in f["chroms"]
    assert "length" in f["chroms"]

    # chroms with extra column
    b = BytesIO()
    f = h5py.File(b, "r+")
    grp = f.create_group("chroms")
    cooler.create._create.write_chroms(grp, chroms.assign(foo=42), h5opts)
    f.flush()
    assert "foo" in f["chroms"]

    # bins
    b = BytesIO()
    f = h5py.File(b, "r+")
    grp = f.create_group("bins")
    cooler.create._create.write_bins(
        grp, bins.assign(foo=42), ["chr1", "chr2", "chr3"], h5opts
    )
    f.flush()
    assert "chrom" in f["bins"]
    assert "start" in f["bins"]
    assert "end" in f["bins"]
    assert "foo" in f["bins"]

    # bins
    b = BytesIO()
    f = h5py.File(b, "r+")
    grp = f.create_group("pixels")
    cooler.create._create.prepare_pixels(
        grp, len(bins), 1000000, ["count", "foo"], {"foo": float}, h5opts
    )
    f.close()

    # pixels
    from multiprocessing import Lock

    nnz, total = cooler.create._create.write_pixels(
        b,
        "pixels",
        ["count", "foo"],
        (pixels.assign(foo=42.0),),
        h5opts,
        Lock(),
    )
    b.seek(0)
    f = h5py.File(b, "r")
    assert "bin1_id" in f["pixels"]
    assert "bin2_id" in f["pixels"]
    assert "count" in f["pixels"]
    assert "foo" in f["pixels"]
    assert f["pixels/foo"].dtype.kind == "f"
    assert total == 1
    assert nnz == 1


def test_many_contigs():
    chroms = pd.DataFrame(
        {
            "name": [f"scaffold_{i:05}" for i in range(4000)],
            "length": np.full(4000, 20),
        },
        columns=["name", "length"],
    )
    bins = cooler.util.binnify(chroms.set_index("name")["length"], 10)
    h5opts = {"compression": "gzip", "compression_opts": 6}

    # chroms
    b = BytesIO()
    f = h5py.File(b, "w")
    cooler.create._create.write_chroms(f.create_group("chroms"), chroms, h5opts)
    cooler.create._create.write_bins(
        f.create_group("bins"), bins, chroms["name"].values, h5opts
    )
    f.flush()
    assert "enum_path" in f["bins/chrom"].attrs

    # TODO: make this more robust
    # assert f['bins/chrom'].attrs['enum_path'] == '/chroms/name'

    cooler.create._create._rename_chroms(
        f, {f"scaffold_{i:05}": f"contig_{i:05}" for i in range(4000)}, h5opts
    )


def test_create_cooler():
    chromsizes = cooler.util.read_chromsizes(op.join(datadir, "toy.chrom.sizes"))
    bins = cooler.util.binnify(chromsizes, 1)
    pixels = pd.read_csv(
        op.join(datadir, "toy.symm.upper.1.zb.coo"),
        sep="\t",
        names=["bin1_id", "bin2_id", "count"],
    )
    pixels["foo"] = 42.0

    with isolated_filesystem():
        cooler.create.create_cooler(
            "test.cool",
            bins,
            pixels,
            assembly="toy",
            metadata={"hello": "world", "list": [1, 2, 3]},
        )

        cooler.create.create_cooler(
            "test.cool::foo/bar",
            bins,
            pixels,
        )

        cooler.create.create_cooler("test.cool", bins, pixels, symmetric_upper=False)

        cooler.create.create_cooler(
            "test.cool",
            bins,
            pixels,
            columns=["count", "foo"],
            dtypes={"foo": np.float64},
        )

        cooler.create.create_cooler(
            "test.cool",
            bins,
            pixels.to_dict(orient="series"),
        )

        cooler.create.create_cooler(
            "test.cool",
            bins,
            (pixels,),
        )

        cooler.create.create_cooler(
            "test.cool",
            bins,
            (pixels.to_dict(orient="series"),),
        )

        two_piece = (pixels.iloc[: len(pixels) // 2], pixels.iloc[len(pixels) // 2 :])
        cooler.create.create_cooler("test.cool", bins, two_piece, ordered=True)
        cooler.create.create_cooler("test.cool", bins, two_piece[::-1], ordered=False)

        many_piece = tuple(
            pixels.iloc[lo:hi] for lo, hi in cooler.util.partition(0, len(pixels), 5)
        )[::-1]
        cooler.create.create_cooler(
            "test.cool", bins, many_piece, ordered=False, max_merge=10
        )

        with pytest.raises(ValueError):
            cooler.create.create_cooler(
                "test.cool",
                bins,
                pixels,
                columns=["count", "missing"],
            )

        with pytest.raises(ValueError):
            cooler.create.create_cooler(
                "test.cool",
                bins[["start", "end"]],
                pixels,
                columns=["count", "missing"],
            )

        with pytest.raises(ValueError):
            cooler.create.create_cooler(
                "test.cool",
                bins[["start", "end"]],
                pixels,
                h5opts={"shuffuffle": "boing"},
            )


def test_create_cooler_from_dask():
    dd = pytest.importorskip("dask.dataframe")

    chromsizes = cooler.util.read_chromsizes(op.join(datadir, "toy.chrom.sizes"))
    bins = cooler.util.binnify(chromsizes, 1)
    pixels = pd.read_csv(
        op.join(datadir, "toy.symm.upper.1.zb.coo"),
        sep="\t",
        names=["bin1_id", "bin2_id", "count"],
    )
    pixels = dd.from_pandas(pixels, npartitions=10)

    with isolated_filesystem():
        cooler.create.create_cooler("test.cool", bins, pixels, ordered=True)

        # TODO: unordered with dask is broken...
        # cooler.create.create_cooler(
        #     "test.cool",
        #     bins,
        #     pixels,
        #     ordered=False
        # )


@pytest.mark.parametrize(
    "fp", [op.join(datadir, "hg19.GM12878-MboI.matrix.2000kb.cool")]
)
def test_create_scool(fp):
    c = cooler.Cooler(fp)
    # chromsizes = c.chromsizes
    bins = c.bins()[:]
    pixels = c.pixels()[:]

    # random and different content to prove only chrom, start, end is linked
    # and the rest is independent for each cell
    from copy import deepcopy

    bins_cell1 = deepcopy(bins)
    bins_cell2 = deepcopy(bins)
    bins_cell3 = deepcopy(bins)
    bins_cell1["weight"] = np.array([0] * len(bins_cell1["start"]))
    bins_cell2["weight"] = np.array([1] * len(bins_cell1["start"]))
    bins_cell3["weight"] = np.array([2] * len(bins_cell1["start"]))

    bins_cell1["KR"] = np.array([3] * len(bins_cell1["start"]))
    bins_cell2["KR"] = np.array([4] * len(bins_cell1["start"]))
    bins_cell3["KR"] = np.array([5] * len(bins_cell1["start"]))

    name_pixel_dict = {"cell1": pixels, "cell2": pixels, "cell3": pixels}
    name_bins_dict = {"cell1": bins_cell1, "cell2": bins_cell2, "cell3": bins_cell3}

    with isolated_filesystem():
        cooler.create_scool("outfile_test.scool", name_bins_dict, name_pixel_dict)
        content_of_scool = cooler.fileops.list_scool_cells("outfile_test.scool")
        content_expected = ["/cells/cell1", "/cells/cell2", "/cells/cell3"]
        for content in content_expected:
            assert content in content_of_scool

        cooler.create_scool("outfile_test.scool", bins, name_pixel_dict)
        content_of_scool = cooler.fileops.list_scool_cells("outfile_test.scool")
        content_expected = ["/cells/cell1", "/cells/cell2", "/cells/cell3"]
        for content in content_expected:
            assert content in content_of_scool