import os
import pathlib
import subprocess
import sys

os.environ["MUJOCO_GL"] = "egl"
os.environ["CUDA_VISIBLE_DEVICES"] = "5"

from lift3d.helpers.common import Logger

CAMERAS = [
    "corner",
    "corner2",
]


TASKS = [
"assembly", "hand-insert","shelf-place","button-press","drawer-open"
]

"""
Easy: button-press, drawer-open, reach, handle-pull, peg-unplug-side, lever-pull, dial-turn
Medium: hammer, sweep-into, bin-picking, push-wall, box-close
Hard: assembly, hand-insert
Very Hard: shelf-place
"""


"""
Result:

Easy: 
        button-press {baseline: [corner1, 1.0], [corner2, 1.0]; Concept_Attention: [corner1, 1.0], [corner2, 1.0]}
        drawer-open {baseline: [corner1, 1.0], [corner2, 1.0]; Concept_Attention: [corner1, 1.0], [corner2, 1.0]}
        reach {baseline: [corner1, 0.60], [corner2, 0.56]; Concept_Attention: [corner1, 0.80], [corner2, 0.52]}
        handle-pull {baseline: [corner1, 1.0], [corner2, 1.0]; Concept_Attention: [corner1, 1.0], [corner2, 1.0]}
        peg-unplug-side {baseline: [corner1, 0.80], [corner2, 1.0]; Concept_Attention: [corner1, 0.96], [corner2, 1.0]}
        lever-pull {baseline: [corner1, 0.64], [corner2, 1.0]; Concept_Attention: [corner1, 0.92], [corner2, 1.0]}
        dial-turn {baseline: [corner1, 1.0], [corner2, 1.0]; Concept_Attention: [corner1, 1.0], [corner2, 1.0]}

Medium: 
        hammer {baseline: [corner1, 0.96], [corner2, 0.88]; Concept_Attention: [corner1, 0.96], [corner2, 1.0]},
        sweep-into {baseline: [corner1, 0.68], [corner2, 0.60]; Concept_Attention: [corner1, 0.84], [corner2, 0.72]},
        bin-picking {baseline: [corner1, 0.92], [corner2, 0.88]; Concept_Attention: [corner1, 0.96], [corner2, 0.92]},
        push-wall {baseline: [corner1, 0.32], [corner2, 0.52]; Concept_Attention: [corner1, 0.36], [corner2, 0.52]},
        box-close {baseline: [corner1, 1.0], [corner2, 0.96]; Concept_Attention: [corner1, 1.0], [corner2, 1.0]},

Hard:   
        assembly {baseline: [corner1, 1.0], [corner2, 1.0]; Concept_Attention: [corner1, 1.0], [corner2, 1.0]},
        hand-insert {baseline: [corner1, 0.72], [corner2. 0.52]; Concept_Attention: [corner1, 0.72], [corner2, 0.72]},

Very Hard: 
        shelf-place {baseline: [corner1, 0.40], [corner2, 0.44]; "Concept_Attention: [corner1, 0.60], [corner2, 0.48]}",
"""


def main():
    code_root = pathlib.Path(__file__).resolve().parent.parent
    tool_path = code_root / "tools" / "gen_data_metaworld.py"
    for task in TASKS:
        for camera in CAMERAS:
            cmd = [
                "python",
                str(tool_path),
                "--task-name",
                task,
                "--camera-name",
                camera,
                "--image-size",
                str(224),
                "--num-episodes",
                str(10),
                "--save-dir",
                "data/metaworld10",
                "--episode-length",
                str(200),
                "--quiet",
            ]
            Logger.log_info(" ".join(cmd))
            subprocess.run(cmd)


if __name__ == "__main__":
    main()
