#!/usr/bin/env python
from __future__ import division, print_function
import os
import numpy as np
import cv2
import StringIO
from PIL import Image as PILImage
import colour
import PIL
import glob

# ROS imports
from rosbag import Bag
from cv_bridge import CvBridge

# ADAPT imports
from postflight_scripts.image_processing import stretch_image_contrast, \
    calibrate_device_rgb


# ---------------------------- Define Parameters -----------------------------
# Location of the bag to extract data from.
bag_fname_glob = '/home/user/adapt_ws/data/2022_jan_flight/*.bag'
bag_fname_glob = '/host_filesystem//mnt/homenas2/noaa_adapt/adapt_flights/2021-11-04/*.bag'
bag_fname_glob = '/host_filesystem//mnt/homenas2/noaa_adapt/adapt_flights/2022-04-21/*.bag'
bag_fname_glob = '/host_filesystem/mnt/homenas2/noaa_adapt/adapt_flights/fairbanks_202109/*.bag'
bag_fname_glob = '/host_filesystem/mnt/homenas2/noaa_adapt/adapt_flights/2022-04-21/*.bag'

# Color calibration matrix.
color_cal_mat_fname = '/home/user/adapt_ws/data/color_cal_mat.txt'

# Gain to correct vignetting.
gain_fname = '/home/user/adapt_ws/data/vignette_correction.tif'

imagery_topic = '/image_raw'    # Imagery topic.
imu_topic = '/an_device/Imu'            # IMU topic
gps_topic = ' /an_device/NavSatFix'     # Raw GPS topic
odometry_topic = '/an_device/Odometry'  # Filtered Odometry topic
filter_status_topic = '/an_device/FilterStatus'         # Filter status topic
ins_system_status_topic = '/an_device/SystemStatus'     # INS status topic

# Directory to explode all of the data to. If None, a directory located
# alongside the bag will be created.
out_dir = None
# ----------------------------------------------------------------------------


# ----------------------------------------------------------------------------
bag_fnames = glob.glob(bag_fname_glob)

if out_dir is None:
    out_dir = os.path.split(bag_fnames[0])[0]

bridge = CvBridge()


# Color calibration matrix.
color_cal_mat = np.loadtxt(color_cal_mat_fname)

# To be applied in linear RGB.
try:
    gain = np.asarray(PIL.Image.open(gain_fname))
    print('Loading gain from \'%s\'' % gain_fname)
except:
    print('Gain file not found')
    gain = None


topics = set()
for bag_fname in bag_fnames:
    base = os.path.splitext(os.path.split(bag_fname)[1])[0]

    if base == 'nav_data':
        continue

    try:
        with Bag(bag_fname, 'r') as ib:
            base = os.path.splitext(os.path.split(bag_fname)[1])[0]

            if base == 'nav_data':
                continue

            image_dir = '%s/images/%s' % (out_dir, base)
            try:
                os.makedirs(image_dir)
            except (IOError, OSError):
                pass

            print('Processing bag: \'%s\'' % bag_fname)
            for topic, msg, t in ib:
                topics.add(topic)

                if topic == imagery_topic:
                    t = msg.header.stamp.to_sec()
                    if hasattr(msg, 'format'):
                        sio = StringIO.StringIO(msg.data)
                        im = PILImage.open( sio )
                        image = np.array( im )
                    else:
                        image = bridge.imgmsg_to_cv2(msg)
                        #image = cv2.cvtColor(image, cv2.COLOR_BAYER_GB2BGR_VNG)
                        #image = cv2.cvtColor(image, cv2.COLOR_BAYER_GR2BGR_VNG)
                        #image = cv2.cvtColor(image, cv2.COLOR_BAYER_RG2BGR_VNG)

                        # ADAPT
                        image = cv2.cvtColor(image, cv2.COLOR_BAYER_BG2BGR_VNG)[:, :, ::-1]

                    image = calibrate_device_rgb(image, gain, color_cal_mat)

                    frame_id = msg.header.frame_id

                    image = stretch_image_contrast(image, 0.5, tile_grid_size=(16, 16))

                    fname = '%s/%i.jpg' % (image_dir, int(np.round(t*1000000)))
                    print('Saving', fname)
                    cv2.imwrite(fname, image[:, :, ::-1],
                                [int(cv2.IMWRITE_JPEG_QUALITY), 98])
    except ValueError:
        pass


odometry = []
nav_bag_out_fname = '%s/nav_data.bag' % out_dir
with Bag(nav_bag_out_fname, 'w') as ob:
    for bag_fname in bag_fnames:
        base = os.path.splitext(os.path.split(bag_fname)[1])[0]

        if 'nav_data' in base:
            continue

        try:
            with Bag(bag_fname, 'r') as ib:
                print('Processing bag: \'%s\'' % bag_fname)
                for topic, msg, t0 in ib:
                    topics.add(topic)

                    try:
                        t = msg.header.stamp.to_sec()
                    except:
                        t = t0.to_sec()

                    if topic == imagery_topic:
                        continue

                    if topic == odometry_topic:
                            odometry.append([t,
                                             msg.pose.pose.position.y,
                                             msg.pose.pose.position.x,
                                             msg.pose.pose.position.z,
                                             msg.pose.pose.orientation.x,
                                             msg.pose.pose.orientation.y,
                                             msg.pose.pose.orientation.z,
                                             msg.pose.pose.orientation.w])

                    ob.write(topic, msg, t0)
        except Exception as e:
            print(e)


odometry = np.array(odometry)
ind = np.argsort(odometry[:, 0])
odometry = odometry[ind]

# Remove redundant times.
odometry = odometry[np.argsort(odometry[:, 0])]
times = odometry[:, 0]
ind = np.where(np.diff(times) > 0)[0]
odometry = odometry[ind]


header = ('time (s)  latitude (deg)  longitude (deg)  height (m)  quatx  '
          'quaty  quatz  quatw')
odometry_txt_fname = '%s/odometry.txt' % out_dir
np.savetxt(odometry_txt_fname, odometry, header=header)