import marimo

__generated_with = "0.14.17"
app = marimo.App(width="medium")


@app.cell
def _():
    import os
    import subprocess
    import sys
    from pathlib import Path
    import marimo as mo
    return Path, mo, os, subprocess, sys


@app.cell
def _(Path, os):
    ROOT = Path(__file__).resolve().parent
    SCRIPT = ROOT / "scripts" / "decode_reimplementation.py"
    OUTPUT = ROOT / "decode_reimplementation_outputs"
    TRAIN_ENV = os.environ.copy()
    TRAIN_ENV["DECODE_DISABLE_TENSORFLOW"] = "1"
    return OUTPUT, SCRIPT, TRAIN_ENV


@app.cell
def _(mo):
    mode = mo.ui.dropdown(["paper_buildings", "meters"], value="paper_buildings", label="Scope")
    model_case = mo.ui.dropdown(["baselines", "lstm", "cnn", "tcn", "timesnet"], value="baselines", label="Model case")
    run = mo.ui.run_button(label="Train full dataset")
    mo.vstack([mo.md("# DECODE full-data experiments"), mode, model_case, run])
    return mode, model_case, run


@app.cell
def _(SCRIPT, TRAIN_ENV, mode, model_case, run, subprocess, sys):
    if run.value:
        args = [sys.executable, str(SCRIPT), "--mode", mode.value]
        if model_case.value == "baselines":
            args += ["--skip-lstm"]
        else:
            args += ["--dl-models", model_case.value, "--epochs", "20", "--batch-size", "64"]
            if model_case.value == "timesnet":
                args += ["--lookback", "144"]
        completed = subprocess.run(args, check=True, env=TRAIN_ENV, text=True)
    else:
        completed = None
    completed
    return


@app.cell
def _(OUTPUT, mo):
    result_files = sorted(OUTPUT.glob("results_*.csv"))
    mo.md("## Outputs\n" + "\n".join(f"- `{p}`" for p in result_files))
    return


if __name__ == "__main__":
    app.run()