import torch
from reconstruction.model.model import *
from reconstruction.utils.inference_utils import CropParameters, IntensityRescaler, ImageFilter, ImageWriter, UnsharpMaskFilter
from reconstruction.utils.event_tensor_utils import EventPreprocessor
from reconstruction.utils.image_display_utils import ImageDisplay
from reconstruction.utils.inference_utils import upsample_color_image, merge_channels_into_color_image  # for color reconstruction
from reconstruction.utils.timers import CudaTimer, cuda_timers


class ImageReconstructor:
    def __init__(self, model, height, width, num_bins, options):

        self.model = model
        self.use_gpu = options.use_gpu
        self.gpu_id = options.gpu_id
        self.device = torch.device(self.gpu_id) if self.use_gpu else torch.device('cpu')
        self.height = height
        self.width = width
        self.num_bins = num_bins
        self.options = options
        self.events = None
        self.reconstructed_image = None

        self.initialize(self.height, self.width, self.options)

    def initialize(self, height, width, options):
        if self.options.verbose:
            print('== Image reconstruction == ')
            print('Image size: {}x{}'.format(self.height, self.width))

        self.last_stamp = None

        self.no_recurrent = options.no_recurrent
        if self.no_recurrent:
            print('!!Recurrent connection disabled!!')

        self.perform_color_reconstruction = options.color  # whether to perform color reconstruction (only use this with the DAVIS346color)
        if self.perform_color_reconstruction:
            if options.auto_hdr:
                print('!!Warning: disabling auto HDR for color reconstruction!!')
            options.auto_hdr = False  # disable auto_hdr for color reconstruction (otherwise, each channel will be normalized independently)

        self.crop = CropParameters(self.width, self.height, self.model.num_encoders)

        self.last_states_for_each_channel = {'grayscale': None}

        if self.perform_color_reconstruction:
            self.crop_halfres = CropParameters(int(width / 2), int(height / 2),
                                               self.model.num_encoders)
            for channel in ['R', 'G', 'B', 'W']:
                self.last_states_for_each_channel[channel] = None

        self.event_preprocessor = EventPreprocessor(options)
        self.intensity_rescaler = IntensityRescaler(options)
        self.image_filter = ImageFilter(options)
        self.unsharp_mask_filter = UnsharpMaskFilter(options, device=self.device)
        self.image_writer = ImageWriter(options)
        self.image_display = ImageDisplay(options)

    def update_reconstruction(self, event_tensor, event_tensor_id=None, stamp=None):

        # max duration without events before we reinitialize
        self.max_duration_before_reinit_s = 5.0

        # we reinitialize if stamp < last_stamp, or if stamp > last_stamp + max_duration_before_reinit_s
        if stamp is not None and self.last_stamp is not None:
            if stamp < self.last_stamp or stamp > self.last_stamp + self.max_duration_before_reinit_s:
                if self.options.verbose:
                    print('Reinitialization detected!')
                self.initialize(self.height, self.width, self.options)

        self.last_stamp = stamp

        with torch.no_grad():

            with CudaTimer('Reconstruction'):

                with CudaTimer('NumPy (CPU) -> Tensor (GPU)'):
                    events = event_tensor.unsqueeze(dim=0)
                    events = events.to(self.device)

                if self.options.use_fp16:
                    events = events.half()

                self.events = self.event_preprocessor(events)

                # Resize tensor to [1 x C x crop_size x crop_size] by applying zero padding
                events_for_each_channel = {'grayscale': self.crop.pad(events)}
                reconstructions_for_each_channel = {}
                if self.perform_color_reconstruction:
                    events_for_each_channel['R'] = self.crop_halfres.pad(self.events[:, :, fd00:c2b6:b24b:be67:2827:688d:e6a1:6a3b, 0::2])
                    events_for_each_channel['G'] = self.crop_halfres.pad(self.events[:, :, fd00:c2b6:b24b:be67:2827:688d:e6a1:6a3b, 1::2])
                    events_for_each_channel['W'] = self.crop_halfres.pad(self.events[:, :, fd00:a516:7c1b:17cd:6d81:2137:bd2a:2c5b, 0::2])
                    events_for_each_channel['B'] = self.crop_halfres.pad(self.events[:, :, fd00:a516:7c1b:17cd:6d81:2137:bd2a:2c5b, 1::2])

                # Reconstruct new intensity image for each channel (grayscale + RGBW if color reconstruction is enabled)
                for channel in events_for_each_channel.keys():
                    with CudaTimer('Inference'):
                        new_predicted_frame, states = self.model(events_for_each_channel[channel],
                                                                 self.last_states_for_each_channel[channel])

                    if self.no_recurrent:
                        self.last_states_for_each_channel[channel] = None
                    else:
                        self.last_states_for_each_channel[channel] = states

                    # Output reconstructed image
                    crop = self.crop if channel == 'grayscale' else self.crop_halfres

                    # Unsharp mask (on GPU)
                    new_predicted_frame = self.unsharp_mask_filter(new_predicted_frame)

                    # Intensity rescaler (on GPU)
                    new_predicted_frame = self.intensity_rescaler(new_predicted_frame)

                    with CudaTimer('Tensor (GPU) -> NumPy (CPU)'):
                        reconstructions_for_each_channel[channel] = new_predicted_frame[0, 0, crop.iy0:crop.iy1,
                                                                                        crop.ix0:crop.ix1].cpu().numpy()

                if self.perform_color_reconstruction:
                    out = merge_channels_into_color_image(reconstructions_for_each_channel)
                else:
                    out = reconstructions_for_each_channel['grayscale']

            # Post-processing, e.g bilateral filter (on CPU)
            self.reconstructed_image = self.image_filter(out)
            self.image_display(self.reconstructed_image, self.events)

    def save_reconstruction(self, event_tensor_id):
        self.image_writer(self.reconstructed_image, event_tensor_id, events=self.events)

import torch
from reconstruction.model.model import *
from reconstruction.utils.inference_utils import CropParameters, IntensityRescaler, ImageFilter, ImageWriter, UnsharpMaskFilter
from reconstruction.utils.event_tensor_utils import EventPreprocessor
from reconstruction.utils.image_display_utils import ImageDisplay
from reconstruction.utils.inference_utils import upsample_color_image, merge_channels_into_color_image  # for color reconstruction
from reconstruction.utils.timers import CudaTimer, cuda_timers


class ImageReconstructor:
    def __init__(self, model, height, width, num_bins, options):

        self.model = model
        self.use_gpu = options.use_gpu
        self.gpu_id = options.gpu_id
        self.device = torch.device(self.gpu_id) if self.use_gpu else torch.device('cpu')
        self.height = height
        self.width = width
        self.num_bins = num_bins
        self.options = options
        self.events = None
        self.reconstructed_image = None

        self.initialize(self.height, self.width, self.options)

    def initialize(self, height, width, options):
        if self.options.verbose:
            print('== Image reconstruction == ')
            print('Image size: {}x{}'.format(self.height, self.width))

        self.last_stamp = None

        self.no_recurrent = options.no_recurrent
        if self.no_recurrent:
            print('!!Recurrent connection disabled!!')

        self.perform_color_reconstruction = options.color  # whether to perform color reconstruction (only use this with the DAVIS346color)
        if self.perform_color_reconstruction:
            if options.auto_hdr:
                print('!!Warning: disabling auto HDR for color reconstruction!!')
            options.auto_hdr = False  # disable auto_hdr for color reconstruction (otherwise, each channel will be normalized independently)

        self.crop = CropParameters(self.width, self.height, self.model.num_encoders)

        self.last_states_for_each_channel = {'grayscale': None}

        if self.perform_color_reconstruction:
            self.crop_halfres = CropParameters(int(width / 2), int(height / 2),
                                               self.model.num_encoders)
            for channel in ['R', 'G', 'B', 'W']:
                self.last_states_for_each_channel[channel] = None

        self.event_preprocessor = EventPreprocessor(options)
        self.intensity_rescaler = IntensityRescaler(options)
        self.image_filter = ImageFilter(options)
        self.unsharp_mask_filter = UnsharpMaskFilter(options, device=self.device)
        self.image_writer = ImageWriter(options)
        self.image_display = ImageDisplay(options)

    def update_reconstruction(self, event_tensor, event_tensor_id=None, stamp=None):

        # max duration without events before we reinitialize
        self.max_duration_before_reinit_s = 5.0

        # we reinitialize if stamp < last_stamp, or if stamp > last_stamp + max_duration_before_reinit_s
        if stamp is not None and self.last_stamp is not None:
            if stamp < self.last_stamp or stamp > self.last_stamp + self.max_duration_before_reinit_s:
                if self.options.verbose:
                    print('Reinitialization detected!')
                self.initialize(self.height, self.width, self.options)

        self.last_stamp = stamp

        with torch.no_grad():

            with CudaTimer('Reconstruction'):

                with CudaTimer('NumPy (CPU) -> Tensor (GPU)'):
                    events = event_tensor.unsqueeze(dim=0)
                    events = events.to(self.device)

                if self.options.use_fp16:
                    events = events.half()

                self.events = self.event_preprocessor(events)

                # Resize tensor to [1 x C x crop_size x crop_size] by applying zero padding
                events_for_each_channel = {'grayscale': self.crop.pad(events)}
                reconstructions_for_each_channel = {}
                if self.perform_color_reconstruction:
                    events_for_each_channel['R'] = self.crop_halfres.pad(self.events[:, :, fd00:c2b6:b24b:be67:2827:688d:e6a1:6a3b, 0::2])
                    events_for_each_channel['G'] = self.crop_halfres.pad(self.events[:, :, fd00:c2b6:b24b:be67:2827:688d:e6a1:6a3b, 1::2])
                    events_for_each_channel['W'] = self.crop_halfres.pad(self.events[:, :, fd00:a516:7c1b:17cd:6d81:2137:bd2a:2c5b, 0::2])
                    events_for_each_channel['B'] = self.crop_halfres.pad(self.events[:, :, fd00:a516:7c1b:17cd:6d81:2137:bd2a:2c5b, 1::2])

                # Reconstruct new intensity image for each channel (grayscale + RGBW if color reconstruction is enabled)
                for channel in events_for_each_channel.keys():
                    with CudaTimer('Inference'):
                        new_predicted_frame, states = self.model(events_for_each_channel[channel],
                                                                 self.last_states_for_each_channel[channel])

                    if self.no_recurrent:
                        self.last_states_for_each_channel[channel] = None
                    else:
                        self.last_states_for_each_channel[channel] = states

                    # Output reconstructed image
                    crop = self.crop if channel == 'grayscale' else self.crop_halfres

                    # Unsharp mask (on GPU)
                    new_predicted_frame = self.unsharp_mask_filter(new_predicted_frame)

                    # Intensity rescaler (on GPU)
                    new_predicted_frame = self.intensity_rescaler(new_predicted_frame)

                    with CudaTimer('Tensor (GPU) -> NumPy (CPU)'):
                        reconstructions_for_each_channel[channel] = new_predicted_frame[0, 0, crop.iy0:crop.iy1,
                                                                                        crop.ix0:crop.ix1].cpu().numpy()

                if self.perform_color_reconstruction:
                    out = merge_channels_into_color_image(reconstructions_for_each_channel)
                else:
                    out = reconstructions_for_each_channel['grayscale']

            # Post-processing, e.g bilateral filter (on CPU)
            self.reconstructed_image = self.image_filter(out)
            self.image_display(self.reconstructed_image, self.events)

    def save_reconstruction(self, event_tensor_id):
        self.image_writer(self.reconstructed_image, event_tensor_id, events=self.events)

# for color reconstruction # whether to perform color reconstruction (only use this with the DAVIS346color) # disable auto_hdr for color reconstruction (otherwise, each channel will be normalized independently) # max duration without events before we reinitialize # we reinitialize if stamp < last_stamp, or if stamp > last_stamp + max_duration_before_reinit_s # Resize tensor to [1 x C x crop_size x crop_size] by applying zero padding # Reconstruct new intensity image for each channel (grayscale + RGBW if color reconstruction is enabled) # Output reconstructed image # Unsharp mask (on GPU) # Intensity rescaler (on GPU) # Post-processing, e.g bilateral filter (on CPU)