# conda activate py37
# conda install -c blazingsql/label/cuda10.2 -c blazingsql -c rapidsai -c nvidia -c conda-forge -c defaults blazingsql python=3.7

from dask_cuda import LocalCUDACluster
from dask.distributed import Client
from blazingsql import BlazingContext
from colorcet import fire
from datashader import Canvas, transfer_functions as tf, utils


def main():

  cluster = LocalCUDACluster()
  client = Client(cluster)


  bc = BlazingContext(dask_client=client, network_interface='lo')
  bc.s3('bsql', bucket_name='blazingsql-colab')

  files_10 = [f's3://bsql/taxi_parquet/yellow_tripdata_2010-0{i}.parquet' if i < 10 
              else f's3://bsql/taxi_parquet/yellow_tripdata_2010-{i}.parquet' for i in range(1, 13)]

  files_11 = [f's3://bsql/taxi_parquet/yellow_tripdata_2011-0{i}.parquet' if i < 10 
              else f's3://bsql/taxi_parquet/yellow_tripdata_2011-{i}.parquet' for i in range(1, 13)]

  f = files_10 + files_11

  bc.create_table('taxi', f)
  bc.sql('select * from taxi').tail()

  pi = 3.141592653589793
  o_shift = 2 * pi * 6378137 / 2

  query = f'''
        SELECT 
            dropoff_longitude * {o_shift} / 180 AS dropoff_x, 
            LOG10(TAN(((90 + dropoff_latitude) * {pi} / 360))) / {pi} / 180 * {o_shift} / 180 AS dropoff_y
        FROM 
            taxi
            WHERE
                pickup_longitude < -73.75 AND pickup_longitude > -74.15
                AND dropoff_longitude < -73.75 AND dropoff_longitude > -74.15
                AND pickup_latitude > 40.68 AND pickup_latitude < 40.84
                AND dropoff_latitude > 40.68 AND dropoff_latitude < 40.84
                '''


  agg = Canvas().points(bc.sql(query), 'dropoff_x', 'dropoff_y')

  img = tf.set_background(tf.shade(agg, cmap=fire), "black")

  utils.export_image(img=img,filename='fire.png', fmt=".png", background=None)

if __name__ == '__main__':
       main()