from pathlib import Path

from farconf import update_fns_to_cli

from cleanba.launcher import FlamingoRun, group_from_fname, launch_jobs
from cleanba.load_and_eval import LoadAndEvalArgs, default_load_and_eval

runs_to_evaluate = [
    *[
        Path("/training/cleanba/064-ablate-back-fixed/wandb") / dname
        for dname in [
            "run-20240627_064349-m7mai9i8",
            "run-20240627_064351-yjzplhr1",
            "run-20240627_064352-yg05sh6c",
            "run-20240627_064401-u9zeuq1b",
            "run-20240627_064402-2pxyq7pe",
            "run-20240627_064402-6r7s6u21",
            "run-20240627_064402-f5osnuxc",
            "run-20240627_064402-vpcqc6lj",
            "run-20240627_064404-56rxj5an",
            "run-20240627_064404-jh3h6nk8",
            "run-20240627_064405-lzazikkc",
            "run-20240627_064407-kgjcqbpw",
            "run-20240627_064408-kfhwjowl",
            "run-20240627_064413-4ayrkwe1",
            "run-20240627_064413-fuougei1",
            "run-20240627_064413-idu01sdd",
            "run-20240627_064413-jasi92bn",
            "run-20240627_064413-m3cs1npf",
            "run-20240627_064413-robdj0r9",
            "run-20240627_064414-afl5t1jz",
            "run-20240627_064414-xaolan0s",
            "run-20240627_064415-gzt5aoao",
            "run-20240627_064415-vn0zt36i",
            "run-20240627_064416-jo2nnr0k",
        ]
    ],
    *[
        Path("/training/cleanba/064-ablate-back/wandb") / dname
        for dname in [
            "run-20240626_082417-4ptrrzwi",
            "run-20240626_082417-nt25b2jg",
            "run-20240626_082435-7la3c32h",
            "run-20240626_082436-k3qq3290",
            "run-20240626_082437-49zw4m3g",
            "run-20240626_082437-mxezj419",
            "run-20240626_082438-f8e2aeu1",
            "run-20240626_082438-hovrwjgw",
            "run-20240626_114016-zx2p2o1g",
            "run-20240626_114017-h4i9lfxd",
            "run-20240626_114033-h7eqkpop",
            "run-20240626_114033-o2tyul1k",
            "run-20240626_114033-okjscwoe",
            "run-20240626_114033-uuilamea",
            "run-20240626_114033-znbc4rz4",
            "run-20240626_114034-gcp8vsgy",
            "run-20240626_114315-tq078cjk",
            "run-20240626_114316-2d1epob7",
            "run-20240626_114327-iwpi0y9h",
            "run-20240626_114334-1naw70o5",
            "run-20240626_114334-pl57snnp",
            "run-20240626_114334-r4x8b25u",
            "run-20240626_114336-47yhch6e",
            "run-20240626_114525-mb3kvc25",
            "run-20240626_114602-6iufrw5c",
            "run-20240626_120531-6rgxmgr8",
            "run-20240626_120602-eql283jg",
            "run-20240626_120612-cdlleqr1",
            "run-20240626_120701-r2kg7ft1",
            "run-20240626_120740-1eqpzsxr",
            "run-20240626_120809-yr6do0ok",
            "run-20240626_120909-59wybefh",
            "run-20240626_122615-vtlw960a",
            "run-20240626_122859-kidkvo5x",
            "run-20240626_123059-0nec5vv9",
            "run-20240626_123150-tjx2sror",
            "run-20240626_123155-1zz5ttlf",
            "run-20240626_125022-z3jeat5e",
            "run-20240626_125558-5jdmxpak",
            "run-20240626_131506-ujsk6uqo",
            "run-20240626_144616-crv0f89c",
            "run-20240626_145037-vmko3gtx",
            "run-20240626_151328-19gx02cv",
            "run-20240626_152313-n7x7v5za",
            "run-20240626_152607-4kexyqew",
            "run-20240626_153449-p1cxaoku",
            "run-20240626_153450-dh3egbev",
            "run-20240626_153457-kxyjdn28",
            "run-20240626_153457-nl55d5fn",
            "run-20240626_153457-v1n55rpl",
            "run-20240626_153458-2rlqdr5c",
            "run-20240626_153458-6vl4ozzu",
            "run-20240626_153459-rhqbpcxb",
            "run-20240626_153459-zfmtq8hv",
            "run-20240626_153500-bdbt5m7j",
            "run-20240626_153500-m364rrjc",
            "run-20240626_153500-ubvdi939",
            "run-20240626_153501-1alsk9oe",
            "run-20240626_153501-yw0s4lp6",
            "run-20240626_153502-hy2t0ngm",
            "run-20240626_153503-4h038ciq",
            "run-20240626_153503-nsq8e8yg",
            "run-20240626_154851-55ey3fe3",
            "run-20240626_154900-9v4g706b",
            "run-20240626_154905-p549yoh9",
            "run-20240626_154932-yg5sox60",
            "run-20240626_154939-g7clt7mu",
            "run-20240626_154941-n5vq7gda",
            "run-20240626_154943-252oepjl",
            "run-20240626_154955-bbl7763z",
            "run-20240626_154955-xmuqzu55",
            "run-20240626_154958-9k6msid4",
            "run-20240626_155001-3bitm9nu",
            "run-20240626_155003-y4z9wxwy",
            "run-20240626_155007-8nbg3jqe",
            "run-20240626_155019-kx0h1axm",
            "run-20240626_155120-nqroy8ks",
            "run-20240626_155656-3p3ztmva",
            "run-20240626_155656-fxuqx7wl",
            "run-20240626_155703-x11a3v75",
            "run-20240626_160305-dqsaq1tb",
            "run-20240626_160309-dx9k0779",
            "run-20240626_160326-1u8j4p16",
            "run-20240626_160401-fqavmnu5",
            "run-20240626_160402-jxkki05l",
            "run-20240626_210139-q7ieg74g",
            "run-20240626_210440-wkaovk0g",
            "run-20240626_211631-4ka49qcg",
            "run-20240626_211632-6nyqqtqc",
            "run-20240626_211632-bpip5wdr",
            "run-20240626_211632-igwodlho",
            "run-20240626_211632-po7xj6iq",
            "run-20240626_211633-e7xpv0t9",
            "run-20240626_211633-kcttmt5u",
            "run-20240626_211633-psqylhm6",
            "run-20240626_211634-1fg88mi7",
            "run-20240626_211634-7wzo7tqn",
            "run-20240626_211634-n2po9vs5",
            "run-20240626_211634-qlcpg4t1",
            "run-20240626_211635-t3p312lr",
            "run-20240626_211635-xec6frqv",
            "run-20240626_211636-kig9w1k6",
            "run-20240626_211636-y9e5e682",
            "run-20240626_211636-yclelqia",
            "run-20240626_211637-1d6jq3uo",
            "run-20240626_211637-kgj9zene",
            "run-20240626_211640-smu1zwej",
            "run-20240626_211643-3nm1b3ys",
            "run-20240626_211643-6tzv00le",
            "run-20240626_211643-wq0srea7",
            "run-20240626_211645-xvyvfflr",
            "run-20240626_211646-hljo6n14",
            "run-20240626_211648-k4cpwm4s",
            "run-20240626_211654-4znyk009",
            "run-20240626_211701-9n690m24",
            "run-20240626_230751-tjblfwv4",
            "run-20240626_230753-0z01y4qe",
            "run-20240626_230753-uyj3x15e",
            "run-20240627_010444-b38z68c3",
            "run-20240627_010444-cdsg92aa",
            "run-20240627_010444-narse2uo",
            "run-20240627_010445-9fgcybke",
            "run-20240627_015705-70d5fj2a",
            "run-20240627_015706-f3gilr54",
            "run-20240627_015707-ew6mcjox",
            "run-20240627_015708-q70lwpqf",
            "run-20240627_020640-rtab9pvw",
            "run-20240627_020641-cptzoqjt",
            "run-20240627_020641-j4nj4taa",
            "run-20240627_020643-gpl1d64b",
        ]
    ],
]


clis: list[list[str]] = []
for load_path in runs_to_evaluate:

    def update(config: LoadAndEvalArgs) -> LoadAndEvalArgs:
        config.load_other_run = load_path
        return config

    cli, _ = update_fns_to_cli(default_load_and_eval, update)
    clis.append(cli)


runs: list[FlamingoRun] = []
RUNS_PER_MACHINE = 2
for update_fns_i in range(0, len(clis), RUNS_PER_MACHINE):
    this_run_clis = [
        ["python", "-m", "cleanba.load_and_eval", *clis[update_fns_i + j]]
        for j in range(min(RUNS_PER_MACHINE, len(clis) - update_fns_i))
    ]
    print(this_run_clis)
    runs.append(
        FlamingoRun(
            this_run_clis,
            CONTAINER_TAG="cbd47ce-main",
            CPU=12,
            MEMORY="40G",
            GPU=1,
            PRIORITY="normal-batch",
            XLA_PYTHON_CLIENT_MEM_FRACTION='".45"',
        )
    )


GROUP: str = group_from_fname(__file__)

if __name__ == "__main__":
    launch_jobs(
        runs,
        group=GROUP,
        job_template_path=Path(__file__).parent.parent.parent / "k8s/runner.yaml",
        project="cleanba",
    )
