#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 20 11:38:08 2021

@author: danapalgokulesh
"""
import numpy as np
import os 
import cv2
img_size = 640
#%%bev-----------
def load_image(path):
    # loads 1 image from dataset, returns img, original hw, resized hw
    img = np.zeros((img_size,img_size,3))
    pts = np.load(path.replace('images','radar_bev').replace('.png','.npy'))
    #pts = pts[[0,1,2,6,9,10,3],:]
    pts[1,:] = (pts[1,:]*3)+(img_size/2); pts[0,:] = pts[0,:]*3
    #pts[0,:] = pts[0,:]/1600*self.img_size; pts[1,:] = pts[1,:]/900*360
    pts[3,:] = (pts[3,:]+5)/57*255; pts[4,:] = (pts[4,:]+98)/169*255; pts[5,:] = (pts[5,:]+51)/109*255;#pts[6,:] = pts[6,:]*255/260
    pts = pts[:,pts[0,:]<=img_size-1]
    pts = pts[:,pts[1,:]<=img_size-1]
    pts = pts[:,pts[0,:]>0]
    pts = pts[:,pts[1,:]>0]
    img[pts[0,:].astype(np.uint),pts[1,:].astype(np.uint),:] = pts.T[:,3:]
    return img, pts

root = '/home/danapalgokulesh/dataset/nuscenes/radar_bev'
troot ='/home/danapalgokulesh/dataset/nuscenes/radar_bev_image'
inputs = os.listdir(root)
for i in inputs[:1]:
    img, pts = load_image(os.path.join(root,i))
    cv2.imwrite(os.path.join(troot,i.replace('npy','png')),img)
    
#%%camera

def load_image(path):
    # loads 1 image from dataset, returns img, original hw, resized hw
    img = np.zeros((img_size,img_size,4))
    pts = np.load(path.replace('images','radar_bev').replace('.png','.npy'))
    pts = pts[[0,1,2,6,9,10,3],:]
    #pts[1,:] = (pts[1,:]*3)+(img_size/2); pts[0,:] = pts[0,:]*3
    pts[0,:] = pts[0,:]/1600*img_size; pts[1,:] = (pts[1,:]/900*360) + 140
    pts[3,:] = (pts[3,:]+5)/57*255; pts[4,:] = (pts[4,:]+98)/169*255; pts[5,:] = (pts[5,:]+51)/109*255;pts[6,:] = pts[6,:]*255/260
    pts = pts[:,pts[0,:]<=img_size-1]
    pts = pts[:,pts[1,:]<=img_size-1]
    pts = pts[:,pts[0,:]>0]
    pts = pts[:,pts[1,:]>0]
    img[pts[1,:].astype(np.uint),pts[0,:].astype(np.uint),:] = pts.T[:,3:]
    return img, pts

root = '/home/danapalgokulesh/dataset/nuscenes/radar_cam'
troot1 ='/home/danapalgokulesh/dataset/nuscenes/radar_cam_image1'
troot2 ='/home/danapalgokulesh/dataset/nuscenes/radar_cam_image2'
inputs = os.listdir(root)
for i in inputs:
    img, pts = load_image(os.path.join(root,i))
    cv2.imwrite(os.path.join(troot1,i.replace('npy','png')),img[:,:,:3])
    cv2.imwrite(os.path.join(troot2,i.replace('npy','png')),img[:,:,-1])