from gfa_core.rnaseq_de_design import plan_rnaseq_de_design


def make_samples(genotypes=("G1", "G2"), times=("0h", "6h", "24h"), reps=2):
    samples = {}
    for genotype in genotypes:
        for time in times:
            for rep in range(1, reps + 1):
                sample = f"{genotype}_{time}_r{rep}"
                samples[sample] = {
                    "condition": "baseline" if time == "0h" else "treatment",
                    "genotype": genotype,
                    "time": time,
                    "replicate": rep,
                    "group_role": "control" if time == "0h" else "compare",
                    "include": True,
                }
    return samples


def test_genotype_time_routes_to_m3_with_per_genotype_baseline_comparisons():
    samples = make_samples()
    plan = plan_rnaseq_de_design(samples, matrix_samples=list(samples))

    assert plan["module"] == "M3"
    assert plan["design"] == "stratified_time"
    assert plan["comparison_axis"] == "time"
    assert plan["stratify_axis"] == "genotype"
    assert plan["reference_time"] == "0h"
    assert plan["compare"] == [
        ["G1__6h", "G1__0h"],
        ["G1__24h", "G1__0h"],
        ["G2__6h", "G2__0h"],
        ["G2__24h", "G2__0h"],
    ]


def test_genotype_time_does_not_require_distinct_condition_labels():
    samples = make_samples()
    for meta in samples.values():
        meta["condition"] = "salt_time_course"

    plan = plan_rnaseq_de_design(samples, matrix_samples=list(samples))

    assert plan["module"] == "M3"
    assert len(plan["compare"]) == 4