# Setup the spark client
from pyspark.sql import SparkSession
warehouseLocation = "/wh"
spark = SparkSession\
	.builder.appName("demo")\
	.config("spark.sql.warehouse.dir", warehouseLocation)\
	.enableHiveSupport()\
	.getOrCreate()
    
# Create test Spark DataFrame
from pyspark.sql.functions import rand
df = spark.range(1 << 22).toDF("id").withColumn("id", rand())
df.printSchema()

# Benchmark time
%time pdf = df.toPandas()
spark.conf.set("spark.sql.execution.arrow.enabled", "true")
%time pdf = df.toPandas()
pdf.describe()

# Do the imports
import os
import ...


# Enable Arrow support.
spark.conf.set("spark.sql.execution.arrow.enabled", "true")
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "2048")

# Enable CUDA
cuda = True

use_cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

bc_model_state = sc.broadcast(models.resnet50(pretrained=True).state_dict())

def get_model_for_eval():
  """Gets the broadcasted model."""
  model = models.resnet50(pretrained=True)
  model.load_state_dict(bc_model_state.value)
  model.eval()
  return model

files_df = spark.createDataFrame(
  map(lambda path: (path,), files), ["path"]
).repartition(10)
display(files_df.limit(10))

class ImageDataset(Dataset):
  def __init__(self, paths, transform=None):
    self.paths = paths
    self.transform = transform
  def __len__(self):
    return len(self.paths)
  def __getitem__(self, index):
    image = default_loader(self.paths[index])
    if self.transform is not None:
      image = self.transform(image)
    return image
    

def predict_batch(paths):
  transform = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225])
  ])
  images = ImageDataset(paths, transform=transform)
  loader = torch.utils.data.DataLoader(images, batch_size=500, num_workers=8)
  model = get_model_for_eval()
  model.to(device)
  all_predictions = []
  with torch.no_grad():
    for batch in loader:
      predictions = list(model(batch.to(device)).cpu().numpy())
      for prediction in predictions:
        all_predictions.append(prediction)
  return pd.Series(all_predictions)

# Set your model
predict_udf = pandas_udf(ArrayType(FloatType()),
			PandasUDFType.SCALAR)(predict_batch)

# Make predictions.
predictions_df = files_df.select(col('path'), 
			predict_udf(col('path')).alias("prediction"))
predictions_df.write.mode("overwrite").parquet(output_file_path)

# Load and check the prediction results.
result_df = spark.read.load(output_file_path)
display(result_df)