import random
import time

import numpy as np
import trackio

import ray
from ray import tune
from ray.air.integrations.trackio import (
    TrackioLoggerCallback,
    setup_trackio,
)
from ray.train import RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer


PROJECT_NAME = "trackio-ray-demo"

HF_DATASET_ID = (
    "AINovice2005/ray-trackio-experiments"  
)
HF_SPACE_ID = (
    "AINovice2005/ray-trackio-dashboard" 
)

NUM_STEPS = 15


def tune_trainable(config):

    for step in range(NUM_STEPS):

        loss = (config["lr"] * 10) / (step + 1) + random.random()
        accuracy = 1 / (loss + 1e-3)

        # Example artifact
        image = np.random.rand(64, 64, 3)

        tune.report(
            {
                "loss": loss,
                "accuracy": accuracy,
                "image_mean": float(image.mean()),
                "step": step,
            }
        )

        time.sleep(0.2)


def run_tune_example():

    tuner = tune.Tuner(
        tune_trainable,
        param_space={
            "lr": tune.grid_search([0.001, 0.01, 0.1]),
        },
        run_config=tune.RunConfig(
            name="trackio-ray-tune-demo",
            callbacks=[
                TrackioLoggerCallback(
                    project=PROJECT_NAME,
                    # Trackio capabilities
                    auto_log_gpu=True,
                    gpu_log_interval=5,
                    dataset_id=HF_DATASET_ID,
                    space_id=HF_SPACE_ID,
                )
            ],
        ),
    )

    results = tuner.fit()

    print("\nTune finished\n")
    print(results)


def train_loop(config):

    run = setup_trackio(
        config=config,
        project=PROJECT_NAME,
        auto_log_gpu=True,
        gpu_log_interval=5,
        dataset_id=HF_DATASET_ID,
        space_id=HF_SPACE_ID,
    )

    for step in range(NUM_STEPS):

        loss = 5 / (step + 1) + random.random()
        throughput = random.uniform(50, 150)

        # Log metrics
        if run:
            run.log(
                {
                    "loss": loss,
                    "throughput": throughput,
                    "step": step,
                },
                step=step,
            )

        # Example artifact
        sample_image = np.random.rand(64, 64, 3)

        if run:
            run.log(
                {
                    "image_mean": float(sample_image.mean()),
                    "image_std": float(sample_image.std()),
                }
            )

        time.sleep(0.2)

    if run:
        run.finish()


def run_train_example():

    trainer = TorchTrainer(
        train_loop_per_worker=train_loop,
        train_loop_config={"lr": 0.01},
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(name="trackio-ray-train-demo"),
    )

    trainer.fit()


def launch_dashboard():

    print("\nLaunching Trackio dashboard...\n")

    trackio.show(
        project=PROJECT_NAME,
        open_browser=True,
    )


if __name__ == "__main__":

    ray.init()

    print("\nRunning Ray Tune experiment\n")
    run_tune_example()

    print("\nRunning Ray Train experiment\n")
    run_train_example()

    print("\nOpening dashboard\n")
    print("Run dashboard manually with:")
    print('trackio show --project "trackio-ray-demo"')

    trackio.finish()
    ray.shutdown()

    print("\nDemo completed\n")