from __future__ import annotations

import numpy as np
import torch
from torchvision import transforms
from torchvision.models import ResNet18_Weights, resnet18
import time
import ray

import daft
from daft import col


import logging
import os

import hydra
import ray
from omegaconf import DictConfig
from data_pipelines.common.logging import setup_logger


# # Wait for Ray cluster to be ready
# @ray.remote
# def warmup():
#     pass
# ray.get([warmup.remote() for _ in range(NUM_GPU_NODES * 23)])


weights = ResNet18_Weights.DEFAULT
transform = transforms.Compose([transforms.ToTensor(), weights.transforms()])


class ResNetModel:
    def __init__(self):
        self.weights = weights
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Device: {self.device}")
        self.model = resnet18(weights=weights).to(self.device)
        self.model.eval()

    @daft.method.batch(
        return_dtype=daft.DataType.string(),
        batch_size=10,
    )
    def __call__(self, images):
        if len(images) == 0:
            return []
        torch_batch = torch.from_numpy(np.array(images.to_pylist())).to(self.device)
        with torch.inference_mode():
            prediction = self.model(torch_batch)
            predicted_classes = prediction.argmax(dim=1).detach().cpu()
            predicted_labels = [self.weights.meta["categories"][i] for i in predicted_classes]
            return predicted_labels


def daft_benchmark(cfg: DictConfig, logger: logging.Logger, **kwargs) -> None:
    """Build eye-gaze action dataset: query Databricks, join with S3 images/labels, write Parquet.

    Reads actions from Databricks, expands to per-frame rows, resolves image and label URLs
    via Daft UDFs, filters to labeled frames, and writes action and image Parquet tables.

    Args:
        cfg: Hydra config with pipeline, env_vars, databricks, ray, logging.
        logger: Logger for progress and debug.
        **kwargs: Additional options passed through.

    Examples:
        Called by Hydra from CLI:
            python data_pipelines/run_pipeline.py pipeline=dms/daft_benchmark
        Or from code with a DictConfig and logger:
            daft_benchmark(cfg, logger)
    """
    model = daft.cls(ResNetModel, gpus=1, max_concurrency=cfg.ray.num_gpus)
    start_time = time.time()
    <http://logger.info|logger.info>(f"Reading S3 Images from from {cfg.pipeline.input_path}")
    df = daft.from_glob_path(f"{cfg.pipeline.input_path}/**/*.png")
    df = df.with_column_renamed("path", "image_url")
    df = df.collect()
    <http://logger.info|logger.info>(f"Number of Images: {df.count_rows()}")
    df = df.with_column(
        "decoded_image",
        df["image_url"].download().decode_image(mode=daft.ImageMode.RGB),
    )
    df = df.with_column(
        "norm_image",
        df["decoded_image"].apply(
            func=lambda image: transform(image),
            return_dtype=daft.DataType.tensor(dtype=daft.DataType.float32(), shape=tuple(cfg.pipeline.image_dim)),
        ),
    )
    df = df.with_column("label", model()(col("norm_image")))
    df = df.select("image_url", "label")
    <http://logger.info|logger.info>(f"Schema: {df.schema()}")
    df.write_parquet(f"{cfg.pipeline.output_path}", write_mode="overwrite")
    end_time = time.time()
    <http://logger.info|logger.info>(f"Time taken: {end_time - start_time} seconds")
    <http://logger.info|logger.info>(f"Images per second: {df.count_rows() / (end_time - start_time)}")
    <http://logger.info|logger.info>(f"Parquet written to {cfg.pipeline.output_path}")

df = daft.from_glob_path(f"{cfg.pipeline.input_path}/**/*.png")

daft.from_glob_path