import os.path as op

import numpy as np
import pandas as pd
from click.testing import CliRunner

import cooler

### COMPUTE ###
from cooler.cli.balance import balance
from cooler.cli.coarsen import coarsen

## REDUCTION ###
from cooler.cli.merge import merge
from cooler.cli.zoomify import zoomify

# import pytest


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


def test_merge():
    runner = CliRunner()
    with runner.isolated_filesystem():
        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        result = runner.invoke(
            merge, ["toy.2.double.cool", f_in, f_in, "--field", "count:dtype=int"]
        )
        assert result.exit_code == 0
        total1 = cooler.Cooler(f_in).pixels()["count"][:].sum()
        total2 = cooler.Cooler("toy.2.double.cool").pixels()["count"][:].sum()
        assert total2 == 2 * total1


def test_coarsen():
    runner = CliRunner()
    with runner.isolated_filesystem():
        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        f_ref = op.join(datadir, "toy.symm.upper.4.cool")
        result = runner.invoke(
            coarsen,
            [f_in, "--factor", "2", "--nproc", "2", "-o", "toy.2.coarsen_2.cool"],
        )
        assert result.exit_code == 0
        pix1 = cooler.Cooler(f_ref).pixels()["count"][:]
        pix2 = cooler.Cooler("toy.2.coarsen_2.cool").pixels()["count"][:]
        assert np.allclose(pix1, pix2)

        result = runner.invoke(
            coarsen,
            [
                f_in,
                "--factor",
                "2",
                "--field",
                "count:dtype=float,agg=mean",
                "-o",
                "toy.2.coarsen_2_mean.cool",
            ],
        )
        assert result.exit_code == 0
        pix2 = cooler.Cooler("toy.2.coarsen_2_mean.cool").pixels()["count"][:]
        assert pix2.dtype.kind == "f"


def test_zoomify():
    runner = CliRunner()

    with runner.isolated_filesystem():
        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        result = runner.invoke(
            zoomify, [f_in, "--balance", "--legacy", "-o", "toy.2.mcool"]
        )
        assert result.exit_code == 0

        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        result = runner.invoke(
            zoomify, [f_in, "--balance", "--nproc", "2", "-o", "toy.2.mcool"]
        )
        assert result.exit_code == 0

        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        result = runner.invoke(
            zoomify, [f_in, "--balance", "--resolutions", "2,4,8", "-o", "toy.2.mcool"]
        )
        assert result.exit_code == 0

        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        result = runner.invoke(
            zoomify,
            [
                f_in,
                "--balance",
                "--resolutions",
                "2,4,8",
                "--field",
                "count:dtype=float,agg=mean",
                "-o",
                "toy.2.mcool",
            ],
        )
        assert result.exit_code == 0
        # pix1 = cooler.Cooler(f_ref).pixels()['count'][:]
        # pix2 = cooler.Cooler('toy.4.cool').pixels()['count'][:]
        # assert np.allclose(pix1, pix2)


def test_balance():
    runner = CliRunner()
    with runner.isolated_filesystem():
        f_in = op.join(datadir, "toy.symm.upper.2.cool")
        result = runner.invoke(
            balance,
            [
                f_in,
                "--ignore-diags",
                "2",
                "--mad-max",
                "0",
                "--min-nnz",
                "0",
                "--tol",
                "0.05",
                "--nproc",
                "2",
                "--stdout",
            ],
        )
        assert result.exit_code == 0
        assert len(result.output.split("\n")) == 32

        # convergence
        result = runner.invoke(
            balance,
            [
                f_in,
                "--ignore-diags",
                "2",
                "--mad-max",
                "0",
                "--min-nnz",
                "0",
                "--tol",
                "0.05",
                "--max-iters",
                "1",
                "--convergence-policy",
                "store_final",
                "--stdout",
            ],
        )
        assert result.exit_code == 0
        assert len(result.output)
        result = runner.invoke(
            balance,
            [
                f_in,
                "--ignore-diags",
                "2",
                "--mad-max",
                "0",
                "--min-nnz",
                "0",
                "--tol",
                "0.05",
                "--max-iters",
                "1",
                "--convergence-policy",
                "discard",
                "--stdout",
            ],
        )
        assert result.exit_code == 0
        assert not result.output
        result = runner.invoke(
            balance,
            [
                f_in,
                "--ignore-diags",
                "2",
                "--mad-max",
                "0",
                "--min-nnz",
                "0",
                "--tol",
                "0.05",
                "--max-iters",
                "1",
                "--convergence-policy",
                "error",
                "--stdout",
            ],
        )
        assert result.exit_code == 1

        # file is unbalanced
        result = runner.invoke(
            balance,
            [f_in, "--check"],
        )
        assert result.exit_code == 1

        # blacklisting regions
        blacklist = pd.DataFrame(
            {
                "chrom": ["chr1", "chr2"],
                "start": [5, 10],
                "end": [10, 20],
            },
            columns=["chrom", "start", "end"],
        )
        blacklist.to_csv("blacklist.bed", sep="\t", index=False, header=False)
        result = runner.invoke(
            balance,
            [
                f_in,
                "--ignore-diags",
                "2",
                "--mad-max",
                "0",
                "--min-nnz",
                "0",
                "--tol",
                "0.05",
                "--blacklist",
                "blacklist.bed",
                "--stdout",
            ],
        )
        assert result.exit_code == 0