import unittest

import bitarray as ba
import numpy as np
import pytest

import ldsc.ldscore as ld
import ldsc.parse as ps


def test_getBlockLefts():
    left = [
        (np.arange(1, 6), 5, np.zeros(5)),
        (np.arange(1, 6), 0, np.arange(0, 5)),
        ((1, 4, 6, 7, 7, 8), 2, (0, 1, 1, 2, 2, 2)),
    ]
    for coords, max_dist, correct in left:
        assert np.all(ld.getBlockLefts(coords, max_dist) == correct)


def test_block_left_to_right():
    left = [
        ((0, 0, 0, 0, 0), (5, 5, 5, 5, 5)),
        ((0, 1, 2, 3, 4, 5), (1, 2, 3, 4, 5, 6)),
        ((0, 0, 2, 2), (2, 2, 4, 4)),
    ]
    for block_left, correct_answer in left:
        block_right = ld.block_left_to_right(block_left)
        assert np.all(block_right == correct_answer)


class test_bed(unittest.TestCase):
    def setUp(self):
        self.M = 8
        self.N = 5
        self.bim = ps.PlinkBIMFile("test/plink_test/plink.bim")

    def test_bed(self):
        bed = ld.PlinkBEDFile("test/plink_test/plink.bed", self.N, self.bim)
        # remove three monomorphic SNPs
        print(bed.geno)
        print(bed.m)
        assert bed.m == 4
        # no individuals removed
        print(bed.n)
        assert bed.n == self.N
        # 5 indivs * 4 polymorphic SNPs
        print(len(bed.geno))
        assert len(bed.geno) == 64
        print(bed.freq)
        correct = np.array([0.59999999999999998, 0.59999999999999998, 0.625, 0.625])
        assert np.all(bed.freq == correct)

    def test_filter_snps(self):
        keep_snps = [1, 4]
        bed = ld.PlinkBEDFile(
            "test/plink_test/plink.bed", self.N, self.bim, keep_snps=keep_snps
        )
        assert bed.m == 1
        assert bed.n == 5
        # pad bits are initialized with random memory --> can't test them
        assert bed.geno[0:10] == ba.bitarray("0001011111")

    def test_filter_indivs(self):
        keep_indivs = [0, 1]
        bed = ld.PlinkBEDFile(
            "test/plink_test/plink.bed", self.N, self.bim, keep_indivs=keep_indivs
        )
        assert bed.m == 2
        assert bed.n == 2
        # pad bits are initialized with random memory --> can't test them
        assert bed.geno[0:4] == ba.bitarray("0001")
        assert bed.geno[8:12] == ba.bitarray("0001")

    def test_filter_indivs_and_snps(self):
        keep_indivs = [0, 1]
        keep_snps = [1, 5]
        bed = ld.PlinkBEDFile(
            "test/plink_test/plink.bed",
            self.N,
            self.bim,
            keep_snps=keep_snps,
            keep_indivs=keep_indivs,
        )
        assert bed.m == 1
        assert bed.n == 2
        print(bed.geno)
        assert bed.geno[0:4] == ba.bitarray("0001")

    def test_bad_filename(self):
        with pytest.raises(ValueError):
            ld.PlinkBEDFile("test/plink_test/plink.bim", 9, self.bim)

    def test_nextSNPs_errors1(self):
        with pytest.raises(ValueError):
            bed = ld.PlinkBEDFile("test/plink_test/plink.bed", self.N, self.bim)
            bed.nextSNPs(0)

    def test_nextSNPs_errors2(self):
        with pytest.raises(ValueError):
            bed = ld.PlinkBEDFile("test/plink_test/plink.bed", self.N, self.bim)
            bed.nextSNPs(5)

    def test_nextSNPs(self):
        for b in [1, 2, 3]:
            bed = ld.PlinkBEDFile("test/plink_test/plink.bed", self.N, self.bim)
            x = bed.nextSNPs(b)
            assert x.shape == (5, b)
            assert np.all(np.abs(np.mean(x, axis=0)) < 0.01)
            assert np.all(np.abs(np.std(x, axis=0) - 1) < 0.01)

    def test_nextSNPs_maf_ref(self):
        b = 4
        bed = ld.PlinkBEDFile("test/plink_test/plink.bed", self.N, self.bim)
        x = bed.nextSNPs(b)
        bed._currentSNP -= b
        y = bed.nextSNPs(b, minorRef=True)
        assert np.all(x == -y)
