import unittest

import numpy as np
import pandas as pd
import pytest
from numpy.testing import assert_allclose
from pandas.testing import assert_frame_equal, assert_series_equal

import ldsc.scripts.munge_sumstats as munge


class Mock:
    """
    Dumb object for mocking args and log
    """

    def __init__(self):
        pass

    def log(self, _):
        pass


log = Mock()
args = munge.parser.parse_args("")


class test_p_to_z(unittest.TestCase):
    def setUp(self):
        self.N = pd.Series([1, 2, 3])
        self.P = pd.Series([0.1, 0.1, 0.1])
        self.Z = pd.Series([1.644854, 1.644854, 1.644854])

    def test_p_to_z(self):
        assert_allclose(munge.p_to_z(self.P, self.N), self.Z, atol=1e-5)


class test_check_median(unittest.TestCase):
    def setUp(self):
        self.x = pd.Series([1, 2, 3])

    def test_good_median(self):
        msg = munge.check_median(self.x, 2, 0, "TEST")
        self.assertEqual(msg, "Median value of TEST was 2.0, which seems sensible.")

    def test_bad_median(self):
        with pytest.raises(ValueError):
            munge.check_median(self.x, 0, 0.1, "TEST")


class test_process_n(unittest.TestCase):
    def setUp(self):
        self.dat = pd.DataFrame(["rs1", "rs2", "rs3"], columns=["SNP"])
        self.dat_filtered = pd.DataFrame(["rs2", "rs3"], columns=["SNP"])
        self.dat_filtered["N"] = [1234, 1234.0]
        self.dat_filtered9999 = pd.DataFrame(["rs2", "rs3"], columns=["SNP"])
        self.dat_filtered9999["N"] = [9999, 9999.0]
        self.args = munge.parser.parse_args("")
        # these flags are either re-set in test cases or should be overridden
        self.args.N = 9999.0
        self.args.N_cas = 9999.0
        self.args.N_con = 9999.0
        self.N_const = pd.Series([1234, 1234, 1234.0])
        self.N = pd.Series([1, 1234, 1234.0])
        self.N_filt = pd.Series([1234, 1234.0])
        self.N9999 = pd.Series([9999.0, 9999, 9999])

    def test_n_col(self):
        self.dat["N"] = self.N
        dat = munge.process_n(self.dat, self.args, log)
        print(dat)
        print(self.dat_filtered)
        assert_frame_equal(dat, self.dat_filtered)

    def test_nstudy(self):
        # should filter on NSTUDY if the --N flag is set, but N gets set to
        # 9999
        self.dat["NSTUDY"] = self.N
        dat = munge.process_n(self.dat, self.args, log)
        assert_frame_equal(dat, self.dat_filtered9999)

    def test_n_cas_con_col(self):
        self.dat["N_CAS"] = self.N
        self.dat["N_CON"] = [0.0, 0, 0]
        dat = munge.process_n(self.dat, self.args, log)
        assert_frame_equal(dat, self.dat_filtered)

    def test_n_flag(self):
        self.args.N = 1234.0
        self.args.N_cas = None
        self.args.N_con = None
        dat = munge.process_n(self.dat, self.args, log)
        assert_series_equal(dat.N, self.N_const, check_names=False)

    def test_n_cas_con_flag(self):
        self.args.N = None
        self.args.N_cas = 1000.0
        self.args.N_con = 234.0
        dat = munge.process_n(self.dat, self.args, log)
        assert_series_equal(dat.N, self.N_const, check_names=False)


def test_filter_pvals():
    P = pd.Series([0, 0.1, 1, 2])
    x = munge.filter_pvals(P, log, args)
    assert_series_equal(x, pd.Series([False, True, True, False]))


def test_single_info():
    dat = pd.Series([0.8, 1, 1])
    x = munge.filter_info(dat, log, args)
    assert_series_equal(x, pd.Series([False, True, True]))


def test_multiple_info():
    i1 = pd.Series([0.8, 1, 1])
    i2 = pd.Series([1.01, 0.5, 9])
    dat = pd.concat([i1, i2], axis=1).reset_index(drop=True)
    dat.columns = ["INFO", "INFO"]
    x = munge.filter_info(dat, log, args)
    assert_series_equal(x, pd.Series([True, False, True]))


def test_filter_frq():
    frq = pd.Series([-1, 0, 0.005, 0.4, 0.6, 0.999, 1, 2])
    x = munge.filter_frq(frq, log, args)
    assert_series_equal(
        x, pd.Series([False, False, False, True, True, False, False, False])
    )


def test_filter_alleles():
    a = pd.Series(["AC", "AG", "CA", "CT", "GA", "GT", "TC", "TG", "DI", "AAT", "RA"])
    x = munge.filter_alleles(a)
    y = pd.Series([i < 8 for i in range(11)])
    assert_series_equal(x, y)


class test_allele_merge(unittest.TestCase):
    def setUp(self):
        self.dat = pd.DataFrame(
            np.transpose([["a", "b", "c"], ["A", "T", "C"], ["C", "G", "A"]])
        )
        self.dat.columns = ["SNP", "A1", "A2"]
        self.alleles = pd.DataFrame(
            np.transpose([["a", "extra", "b", "c"], ["AG", "TC", "AC", "AC"]])
        )
        self.alleles.columns = ["SNP", "MA"]

    def test_merge(self):
        x = munge.allele_merge(self.dat, self.alleles, log)
        answer = pd.DataFrame(
            np.transpose(
                [["a", "extra", "b", "c"], ["a", "a", "T", "C"], ["a", "a", "G", "A"]]
            )
        )
        answer.columns = ["SNP", "A1", "A2"]
        answer.loc[[0, 1], ["A1", "A2"]] = float("nan")
        assert_frame_equal(x, answer)


class test_parse_dat(unittest.TestCase):
    def setUp(self):
        dat = pd.DataFrame()
        dat["SNP"] = ["rs" + str(i) for i in range(10)]
        dat["A1"] = ["A" for _ in range(10)]
        dat["A2"] = ["G" for _ in range(10)]
        dat["INFO"] = np.ones(10)
        dat["FRQ"] = np.ones(10) / 2
        dat["P"] = np.ones(10)
        self.dat = dat
        self.dat_gen = [dat.loc[0:4, :], dat.loc[5:9, :].reset_index(drop=True)]
        self.convert_colname = {x: x for x in self.dat_gen[0].columns}
        self.args = munge.parser.parse_args("")

    def tearDown(self):
        self.args = munge.parser.parse_args("")

    def test_no_alleles(self):
        # test that it doesn't crash with no allele columns and the
        # --no-alleles flag set
        dat = self.dat.drop(["A1", "A2"], axis=1)
        dat_gen = [dat.loc[0:4, :], dat.loc[5:9, :].reset_index(drop=True)]
        self.args.no_alleles = True
        dat = munge.parse_dat(dat_gen, self.convert_colname, None, log, self.args)
        assert_frame_equal(dat, self.dat.drop(["INFO", "FRQ", "A1", "A2"], axis=1))

    def test_merge_alleles(self):
        self.args.merge_alleles = True
        merge_alleles = pd.DataFrame()
        merge_alleles["SNP"] = ["rs" + str(i) for i in range(3)]
        merge_alleles["MA"] = ["AG", "AG", "AG"]
        dat = munge.parse_dat(
            self.dat_gen, self.convert_colname, merge_alleles, log, self.args
        )
        print(self.dat.loc[0:2, ["SNP", "A1", "A2", "P"]])
        assert_frame_equal(dat, self.dat.loc[0:2, ["SNP", "A1", "A2", "P"]])

    def test_standard(self):
        dat = munge.parse_dat(self.dat_gen, self.convert_colname, None, log, self.args)
        assert_frame_equal(dat, self.dat.drop(["INFO", "FRQ"], axis=1))

    def test_na(self):
        self.dat.loc[0, "SNP"] = float("NaN")
        self.dat.loc[1, "A2"] = float("NaN")
        self.dat_gen = [
            self.dat.loc[0:4, :],
            self.dat.loc[5:9, :].reset_index(drop=True),
        ]
        dat = munge.parse_dat(self.dat_gen, self.convert_colname, None, log, self.args)
        assert_frame_equal(
            dat, self.dat.loc[2:, ["SNP", "A1", "A2", "P"]].reset_index(drop=True)
        )


def test_clean_header():
    assert munge.clean_header("foo-bar.foo_BaR") == "FOO_BAR_FOO_BAR"


class test_parse_flag_cnames(unittest.TestCase):
    def setUp(self):
        self.args = munge.parser.parse_args("")

    def test_basic(self):
        self.args.nstudy = "nstudy1"
        self.args.snp = "snp1"
        self.args.N_col = "n.col1"
        self.args.N_cas_col = "n-cas.col1"
        self.args.N_con_col = "n-con.col1"
        self.args.a1 = "a11"
        self.args.a2 = "a21"
        self.args.p = "p1"
        self.args.frq = "frq1"
        self.args.info = "info1"
        self.args.info_list = "info111,info222"
        self.args.signed_sumstats = "beta1,0"
        x, y = munge.parse_flag_cnames(log, self.args)
        self.assertEqual(y, 0)
        self.assertEqual(x["NSTUDY1"], "NSTUDY")
        self.assertEqual(x["SNP1"], "SNP")
        self.assertEqual(x["N_COL1"], "N")
        self.assertEqual(x["N_CAS_COL1"], "N_CAS")
        self.assertEqual(x["N_CON_COL1"], "N_CON")
        self.assertEqual(x["A11"], "A1")
        self.assertEqual(x["A21"], "A2")
        self.assertEqual(x["P1"], "P")
        self.assertEqual(x["FRQ1"], "FRQ")
        self.assertEqual(x["INFO1"], "INFO")
        self.assertEqual(x["INFO111"], "INFO")
        self.assertEqual(x["INFO222"], "INFO")

    def test_sign_error(self):
        self.args.signed_sumstats = "1,2,3"
        with pytest.raises(ValueError):
            munge.parse_flag_cnames(log, self.args)
        self.args.signed_sumstats = "BETA,B"
        with pytest.raises(ValueError):
            munge.parse_flag_cnames(log, self.args)
        self.args.signed_sumstats = "BETA"
        with pytest.raises(ValueError):
            munge.parse_flag_cnames(log, self.args)


class test_cname_map(unittest.TestCase):
    def setUp(self):
        pass

    def test_no_flags(self):
        x = munge.get_cname_map({}, munge.default_cnames, [])
        self.assertEqual(x, munge.default_cnames)

    def test_ignore(self):
        ignore = ["sNp", "a1"]
        flag_cnames = {"SNP": "SNP", "ASDF": "ASDF", "N": "FOOBAR"}
        x = munge.get_cname_map(flag_cnames, munge.default_cnames, ignore)
        # check that ignore columns are ignored
        with pytest.raises(KeyError):
            x.__getitem__("SNP")
        with pytest.raises(KeyError):
            x.__getitem__("A1")
        # check that flag columns make it into the dict
        self.assertEqual(x["ASDF"], "ASDF")
        # check that default columns make it into the dict
        self.assertEqual(x["A2"], "A2")
        # check that flags override default
        self.assertEqual(x["N"], "FOOBAR")


class test_end_to_end(unittest.TestCase):
    def setUp(self):
        self.args = munge.parser.parse_args("")
        self.args.sumstats = "test/munge_test/sumstats"
        self.args.out = "asdf"
        self.args.daner = True

    def test_basic(self):
        x = munge.munge_sumstats(self.args, p=False)
        correct = pd.read_csv("test/munge_test/correct.sumstats", sep=r"\s+", header=0)
        assert_frame_equal(x, correct)

    def test_merge_alleles(self):
        self.args.merge_alleles = "test/munge_test/merge_alleles"
        x = munge.munge_sumstats(self.args, p=False)
        correct = pd.read_csv(
            "test/munge_test/correct_merge.sumstats", sep=r"\s+", header=0
        )
        assert_frame_equal(x, correct)

    def test_bad_merge_alleles(self):
        self.args.merge_alleles = "test/munge_test/merge_alleles_bad"
        with pytest.raises(ValueError):
            munge.munge_sumstats(self.args, p=False)

    def test_bad_flags1(self):
        self.args.sumstats = None
        with pytest.raises(ValueError):
            munge.munge_sumstats(self.args, p=False)

    def test_bad_flags2(self):
        self.args.out = None
        with pytest.raises(ValueError):
            munge.munge_sumstats(self.args, p=False)

    def test_bad_flags3(self):
        self.args.merge_alleles = "foo"
        self.args.no_alleles = "bar"
        with pytest.raises(ValueError):
            munge.munge_sumstats(self.args, p=False)

    def test_bad_sumstats1(self):
        self.args.signed_sumstats = "OR,0"
        with pytest.raises(ValueError):
            munge.munge_sumstats(self.args, p=False)

    def test_bad_sumstats2(self):
        self.args.signed_sumstats = "BETA,0"
        with pytest.raises(ValueError):
            munge.munge_sumstats(self.args, p=False)
