from metaflow import (
    FlowSpec,
    step,
    metaflow_ray,
    batch,
    card,
    current,
    environment,
    pypi,
    kubernetes
)

NUM_WORKERS = 1
USE_GPU = True
RESOURCES = dict(cpu=1, memory=16000, node_selector="profile=cpu-ssd")

COMMON_PKGS = {
    "ray[data,train,tune]": "2.39",
    "lightning":"2.1.0",
    "torch":">=2.1.0,<=2.3.1",
    "pandas":"2.1.3",
    "metaflow-ray":"0.1.0"

}

class TorchTrainerGPU(FlowSpec):
    n_cpu = RESOURCES["cpu"]

    @pypi(packages=COMMON_PKGS)
    @step
    def start(self):
        self.next(self.train, num_parallel=NUM_WORKERS)

    @pypi(packages=COMMON_PKGS)
    @metaflow_ray(all_nodes_started_timeout=90)
    @kubernetes(**RESOURCES)
    @card
    @step
    def train(self):
        import ray
        from model import train_loop_per_worker
        from ray.air.config import ScalingConfig
        from ray.air import RunConfig, CheckpointConfig
        from ray.train.torch import TorchTrainer

        context = ray.init()
        print(f"Ray Dashboard: {context.dashboard_url}")
        print("Ray initialized in the %s step." % current.step_name)

        run_id = current.run_id
        step_name = current.step_name
        control_task_id = current.task_id
        # UBF handling for multinode case

        print("Ray nodes: ", ray.nodes())
        # self._control_mapper_tasks = [
        #     "{}/{}/{}".format(run_id, step_name, task_id)
        #     for task_id in [current.task_id]
        #                    + [
        #                        "%s-worker-%d" % (current.task_id, idx)
        #                        for idx in range(NUM_WORKERS - 1)
        #                    ]
        # ]
        print(f"Control Mapper: {self._control_mapper_tasks}")
        print("Ray cluster resources:")
        for k, v in ray.cluster_resources().items():
            if "memory" in k.lower():
                print("%s: %sGB" % (k, round(int(v) / (1024 * 1024 * 1024), 2)))
            else:
                print("%s: %s" % (k, v))

        # Define configurations.
        train_loop_config = {"num_epochs": 20, "lr": 0.01, "batch_size": 32}
        scaling_config = ScalingConfig(num_workers=NUM_WORKERS,resources_per_worker={"CPU":1}, use_gpu=False)
        run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=1))

        # Define datasets.
        train_dataset = ray.data.from_items(
            [{"input": [x], "label": [2 * x + 1]} for x in range(2000)]
        )
        datasets = {"train": train_dataset}

        # Initialize the Trainer.
        trainer = TorchTrainer(
            train_loop_per_worker=train_loop_per_worker,
            train_loop_config=train_loop_config,
            scaling_config=scaling_config,
           run_config=run_config,
            datasets=datasets
        )

        # Train the model.
        result = trainer.fit()
        # Inspect the results.
        final_loss = result.metrics["loss"]
        print(final_loss)
        self.next(self.join)

    @pypi(packages=COMMON_PKGS)
    @step
    def join(self, inputs):
        self.merge_artifacts(inputs)
        self.next(self.end)

    @pypi(packages=COMMON_PKGS)
    @step
    def end(self):
        print(self.result.path)


if __name__ == "__main__":
    TorchTrainerGPU()