import hail as hl


hl.init(
        log="/constraint_pipeline_testing.log",
        tmp_dir="gs://gnomad-tmp-4day",
    )


use_pops=False
pops = constraint_res.POPS if use_pops else ()


from gnomad_constraint.pipeline.constraint_pipeline import get_constraint_resources
from gnomad_constraint.utils.constraint import apply_models

version="2.1.1"
resources = get_constraint_resources(
        version,
        use_v2_release_mutation_ht=False,
        use_v2_release_context_ht=False,
        custom_vep_annotation="transcript_consequences",
        overwrite=False,
        test=False,
    )

res = resources.apply_models
r="autosome_par"


exome_ht = getattr(res, f"preprocessed_{r}_exomes_ht").ht()
context_ht = getattr(res, f"preprocessed_{r}_context_ht").ht()
mutation_ht = res.mutation_ht.ht().select("mu_snp")


# Test changing partitions
#mutation_ht = mutation_ht = mutation_ht.repartition(1)
#mutation_ht = mutation_ht.naive_coalesce(1)

oe_ht = apply_models(
                    exome_ht,
                    context_ht,
                    mutation_ht,
                    getattr(res, f"model_{r}_plateau_ht").he(),
                    getattr(res, f"model_{r}_coverage_ht").he(),
                    max_af=0.001,
                    pops=pops,
                    custom_vep_annotation="transcript_consequences",
                )