import os
import unittest
import pytest
import numpy as np

import hail as hl
import hail.expr.aggregators as agg
import hail.utils as utils
from hail.linalg import BlockMatrix
from hail.utils import FatalError
from ..helpers import (startTestHailContext, stopTestHailContext, resource,
                       skip_unless_spark_backend, fails_local_backend)

setUpModule = startTestHailContext
tearDownModule = stopTestHailContext


class Tests(unittest.TestCase):
    @unittest.skipIf('HAIL_TEST_SKIP_PLINK' in os.environ, 'Skipping tests requiring plink')
    @fails_local_backend()
    def test_impute_sex_same_as_plink(self):
        ds = hl.import_vcf(resource('x-chromosome.vcf'))

        sex = hl.impute_sex(ds.GT, include_par=True)

        vcf_file = utils.uri_path(utils.new_temp_file(prefix="plink", extension="vcf"))
        out_file = utils.uri_path(utils.new_temp_file(prefix="plink"))

        hl.export_vcf(ds, vcf_file)

        utils.run_command(["plink", "--vcf", vcf_file, "--const-fid",
                           "--check-sex", "--silent", "--out", out_file])

        plink_sex = hl.import_table(out_file + '.sexcheck',
                                    delimiter=' +',
                                    types={'SNPSEX': hl.tint32,
                                           'F': hl.tfloat64})
        plink_sex = plink_sex.select('IID', 'SNPSEX', 'F')
        plink_sex = plink_sex.select(
            s=plink_sex.IID,
            is_female=hl.if_else(plink_sex.SNPSEX == 2,
                              True,
                              hl.if_else(plink_sex.SNPSEX == 1,
                                      False,
                                      hl.null(hl.tbool))),
            f_stat=plink_sex.F).key_by('s')

        sex = sex.select('is_female', 'f_stat')

        self.assertTrue(plink_sex._same(sex.select_globals(), tolerance=1e-3))

        ds = ds.annotate_rows(aaf=(agg.call_stats(ds.GT, ds.alleles)).AF[1])

        self.assertTrue(hl.impute_sex(ds.GT)._same(hl.impute_sex(ds.GT, aaf='aaf')))

    backend_name = os.environ.get('HAIL_QUERY_BACKEND', 'spark')
    linreg_functions = [hl.linear_regression_rows, hl._linear_regression_rows_nd] if backend_name == "spark" else [hl._linear_regression_rows_nd]

    def test_linreg_basic(self):
        phenos = hl.import_table(resource('regressionLinear.pheno'),
                                 types={'Pheno': hl.tfloat64},
                                 key='Sample')
        covs = hl.import_table(resource('regressionLinear.cov'),
                               types={'Cov1': hl.tfloat64, 'Cov2': hl.tfloat64},
                               key='Sample')

        mt = hl.import_vcf(resource('regressionLinear.vcf'))
        mt = mt.annotate_cols(pheno=phenos[mt.s].Pheno, cov=covs[mt.s])
        mt = mt.annotate_entries(x=mt.GT.n_alt_alleles()).cache()

        for linreg_function in self.linreg_functions:
            t1 = linreg_function(
                y=mt.pheno, x=mt.GT.n_alt_alleles(), covariates=[1.0, mt.cov.Cov1, mt.cov.Cov2 + 1 - 1])
            t1 = t1.select(p=t1.p_value)

            t2 = linreg_function(
                y=mt.pheno, x=mt.x, covariates=[1.0, mt.cov.Cov1, mt.cov.Cov2])
            t2 = t2.select(p=t2.p_value)

            t3 = linreg_function(
                y=[mt.pheno], x=mt.x, covariates=[1.0, mt.cov.Cov1, mt.cov.Cov2])
            t3 = t3.select(p=t3.p_value[0])

            t4 = linreg_function(
                y=[mt.pheno, mt.pheno], x=mt.x, covariates=[1.0, mt.cov.Cov1, mt.cov.Cov2])
            t4a = t4.select(p=t4.p_value[0])
            t4b = t4.select(p=t4.p_value[1])

            self.assertTrue(t1._same(t2))
            self.assertTrue(t1._same(t3))
            self.assertTrue(t1._same(t4a))
            self.assertTrue(t1._same(t4b))

    def test_linreg_pass_through(self):
        phenos = hl.import_table(resource('regressionLinear.pheno'),
                                 types={'Pheno': hl.tfloat64},
                                 key='Sample')
        covs = hl.import_table(resource('regressionLinear.cov'),
                               types={'Cov1': hl.tfl
oat64, 'Cov2': hl.tfloat64},
                               key='Sample')

        mt = hl.import_vcf(resource('regressionLinear.vcf')).annotate_rows(foo = hl.struct(bar=hl.rand_norm(0, 1)))

        for linreg_function 
in self.linreg_functions:

            # single group
            lr_result = linreg_function(phenos[mt.s].Pheno, mt.GT.n_alt_alleles(), [1.0],
                                        pass_through=['filters', mt.foo.bar, mt.qual])

            assert mt.aggregate_rows(hl.agg.all(mt.foo.bar == lr_result[mt.row_key].bar))

            # chained
            lr_result = linreg_function([[phenos[mt.s].Pheno]], mt.GT.n_alt_alleles(), [1.0],
                                        pass_through=['filters', mt.foo.bar, mt.qual])

            assert mt.aggregate_rows(hl.agg.all(mt.foo.bar == lr_result[mt.row_key].bar))

            # check types
            assert 'filters' in lr_result.row
            assert lr_result.filters.dtype == mt.filters.dtype

            assert 'bar' in lr_result.row
            assert lr_result.bar.dtype == mt.foo.bar.dtype

            assert 'qual' in lr_result.row
            assert lr_result.qual.dtype == mt.qual.dtype

            # should run successfully with key fields
            linreg_function([[phenos[mt.s].Pheno]], mt.GT.n_alt_alleles(), [1.0],
                            pass_through=['locus', 'alleles'])

            # complex expression
            with pytest.raises(ValueError):
                linreg_function([[phenos[mt.s].Pheno]], mt.GT.n_alt_alleles(), [1.0],
                                pass_through=[mt.filters.length()])

    def test_linreg_chained(self):
        phenos = hl.import_table(resource('regressionLinear.pheno'),
                                 types={'Pheno': hl.tfloat64},
                                 key='Sample')
        covs = hl.import_table(resource('regressionLinear.cov'),
                               types={'Cov1': hl.tfloat64, 'Cov2': hl.tfloat64},
                               key='Sample')

        mt = hl.import_vcf(resource('regressionLinear.vcf'))
        mt = mt.annotate_cols(pheno=phenos[mt.s].Pheno, cov=covs[mt.s])
        mt = mt.annotate_entries(x=mt.GT.n_alt_alleles()).cache()

        for linreg_function in self.linreg_functions:

            t1 = linreg_function(y=[[mt.pheno], [mt.pheno]], x=mt.x, covariates=[1, mt.cov.Cov1, mt.cov.Cov2])

            def all_eq(*args):
                pred = True
                for a in args:
                    if isinstance(a, hl.expr.Expression) \
                            and isinstance(a.dtype, hl.tarray) \
                            and isinstance(a.dtype.element_type, hl.tarray):
                        pred = pred & (hl.all(lambda x: x,
                                              hl.map(lambda elt: ((hl.is_nan(elt[0]) & hl.is_nan(elt[1])) | (elt[0] == elt[1])),
                                                     hl.zip(a[0], a[1]))))
                    else:
                        pred = pred & ((hl.is_nan(a[0]) & hl.is_nan(a[1])) | (a[0] == a[1]))
                return pred

            assert t1.aggregate(hl.agg.all(
                all_eq(t1.n,
                       t1.sum_x,
                       t1.y_transpose_x,
                       t1.beta,
                       t1.standard_error,
                       t1.t_stat,
                       t1.p_value)))

            mt2 = mt.filter_cols(mt.cov.Cov2 >= 0)
            mt3 = mt.filter_cols(mt.cov.Cov2 <= 0)

            # test that chained linear regression can replicate separate calls with different missingness
            t2 = hl.linear_regression_rows(y=mt2.pheno, x=mt2.x, covariates=[1, mt2.cov.Cov1])
            t3 = hl.linear_regression_rows(y=mt3.pheno, x=mt3.x, covariates=[1, mt3.cov.Cov1])

            chained = hl.linear_regression_rows(y=[[hl.case().when(mt.cov.Cov2 >= 0, mt.pheno).or_missing()],
                                                   [hl.case().when(mt.cov.Cov2 <= 0, mt.pheno).or_missing()]],
                                            