#!/usr/bin/env python
# coding: utf-8

# # Pandora: a new stereo matching framework
# <img src="img/logo-cnes-triangulaire.jpg" width="200" height="200">
# 
# *Cournet, M., Sarrazin, E., Dumas, L., Michel, J., Guinet, J., Youssefi, D., Defonte, V., Fardet, Q., 2020. Ground-truth generation and disparity estimation for optical satellite imagery. ISPRS - International Archives of the Photogrammetry, Remote Sensing and Spatial Information Sciences.*

# # Introduction and basic usage

# #### Imports and external functions

# In[ ]:


import io
from IPython.display import Image, display
import io
import matplotlib.pyplot as plt
import numpy as np
import os
from pathlib import Path
import rasterio


# In[ ]:


from snippets.utils import *


# In[ ]:


def plot_image(img, title=None, output_dir=None, cmap="viridis"):
    fig = plt.figure()
    plt.title(title)
    plt.imshow(img, cmap=cmap, vmin=np.min(img), vmax=np.max(img))
    plt.colorbar()
    if output_dir is not None:
        fig.savefig(os.path.join(output_dir,title + '.pdf'))


# In[ ]:


def plot_state_machine(machine):
    stream = io.BytesIO()
    try:
        pandora_machine.get_graph().draw(stream, prog='dot', format='png')
        display(Image(stream.getvalue()))
    except:
        print("It is not possible to show the graphic of the state machine. To solve it, please install graphviz on your system (apt-get install graphviz if operating in Linux) and install python package with pip install graphviz")


# # What is Pandora ?
# 
# * Pandora is a Toolbox to estimate disparity 
# * It is inspired by the work of [Scharstein 2002]
# * Pandora embeds several state-of-art algorithms
# * It is easy to configure and modular
# * Will be used in the 3D reconstruction pipeline CARS for CO3D mission 
# 
# [Scharstein 2002] *A Taxonomy and Evaluation of Dense Two-Frame Stereo Correspondence Algorithms*, D. Scharstein and R. Szeliski, 
# vol. 47, International Journal of Computer Vision}, 2002

# <img src="img/logo_diagram.png" width="500">

# ## Inputs
# 
# * Stereo rectified image pair (with associated masks)
# * Disparity range to explore 
# * Configuration file
# 
# ## Outputs
# 
# * Disparity and validity maps in left image geometry 
# * Disparity and validity maps in righ image geometry (*optional*)

# ## Pandora's pipeline
# 
# Pandora provides the following steps:
# * matching cost computation (**mandatory**)
# * cost aggregation
# * cost optimization
# * disparity computation (**mandatory**)
# * subpixel disparity refinement
# * disparity filtering
# * validation
# * multiscale processing

# <img src="img/schema_complet.png" width="1000">

# ### Available implementations for each step
# 
# | Step                      | Algorithms implemented |
# |:--------------------------|:-----------------------|
# | Matching cost computation | Census / SAD / SSD / ZNNC / MC-CNN |
# | Cost aggregation          | Cross Based Cost Aggregation        |
# | Cost optimization         | SGM               |
# | Disparity computation     | Winner-Take-All   |
# | Subpixel disparity refinement | Vfit / Quadratic |
# | Disparity filtering       | Median / Bilateral |
# | Validation                | Cross checking     |
# | Multiscale                | Fixed zoom pyramid |

# # Pandora execution options with state machine

# #### Imports of pandora 

# In[ ]:


# Load pandora imports
import pandora
from pandora.img_tools import read_img
from pandora.check_json import check_pipeline_section, concat_conf, memory_consumption_estimation
from pandora.state_machine import PandoraMachine
from pandora import import_plugin, check_conf


# #### (Optional) If Pandora plugins are to be used, import them

# Available Pandora Plugins include :
# - MC-CNN Matching cost computation
# - SGM Optimization

# In[ ]:


# Load plugins
import_plugin()


# #### Load and visualize input data 

# Provide image path

# In[ ]:


# Paths to left and right images
img_left_path = "data/Cones_LEFT.tif"
img_right_path = "data/Cones_RIGHT.tif"
# Paths to masks (None if not provided)
left_mask_path = None
right_mask_path = None


# Provide image configuration

# In[ ]:


image_cfg = {'image': {'no_data_left': np.nan, 'no_data_right': np.nan}}


# Provide output directory to write results

# In[ ]:


output_dir = os.path.join(os.getcwd(),"output")
# If necessary, create output dir
Path(output_dir).mkdir(exist_ok=True,parents=True)


# Convert input data to dataset

# In[ ]:


img_left = read_img(img_left_path, no_data=image_cfg['image']['no_data_left'],
                       mask=left_mask_path)
img_right = read_img(img_right_path, no_data=image_cfg['image']['no_data_right'],
                       mask=right_mask_path)


# Visualize input data

# In[ ]:


plot_image(img_left.im, "Left input image", output_dir, cmap="gray")


# # Option 1 : trigger all the steps of the machine at ones

# #### Instantiate the machine

# In[ ]:


pandora_machine = PandoraMachine()


# #### Define pipeline configuration

# In[ ]:


user_pipeline_cfg = {
    'pipeline':{
        "right_disp_map": {
            "method": "accurate"
        },
        "matching_cost" : {
            "matching_cost_method": "zncc",
            "window_size": 5,
            "subpix": 4
        },
        "disparity": {
            "disparity_method": "wta",
            "invalid_disparity": "NaN"
        },
        "refinement": {
            "refinement_method": "quadratic"
        },
        "filter": {
            "filter_method": "median"
        },
        "validation": {
            "validation_method": "cross_checking"
        },
        "filter.this_time_after_validation" : {
          "filter_method": "median",
          "filter_size": 3
        }        
    }
}


# Disparity interval used

# In[ ]:


disp_min = -60
disp_max = 0


# #### Check the configuration and sequence of steps

# In[ ]:


checked_cfg = check_pipeline_section(user_pipeline_cfg, pandora_machine)


# In[ ]:


pipeline_cfg = checked_cfg['pipeline']
print(pipeline_cfg)


# #### Estimate the memory consumption of the pipeline

# In[ ]:


min_mem_consump, max_mem_consump = memory_consumption_estimation(user_pipeline_cfg, [img_left_path, disp_min, disp_max], pandora_machine)
print("Estimated maximum memory consumption between {:.2f} GiB and {:.2f} GiB".format(min_mem_consump, max_mem_consump))


# #### Prepare the machine 

# In[ ]:


pandora_machine.run_prepare(pipeline_cfg, img_left, img_right, disp_min, disp_max)


# #### Trigger all the steps of the machine at ones

# In[ ]:


left_disparity, right_disparity = pandora.run(pandora_machine, img_left, img_right, disp_min, disp_max, pipeline_cfg)


# Visualize output disparity map

# In[ ]:


plot_image(left_disparity.disparity_map, "Left disparity map", output_dir, cmap=pandora_cmap())


# # Option 2 : trigger the machine step by step

# The implementation of Pandora with a state machine makes it possible to set up a more flexible pipeline, which makes it possible to choose via a configuration file the steps as well as the order of the steps that one wishes to follow in Pandora.
# 
# Moreover, the state machine allows to run each step of the pipeline independently, giving the possibility to save and visualize the results after each step. 

# The state machine has three states : 
# * Begin
# * Cost volume
# * Disparity map 
# 
# Being the connections between them the different steps of the pipeline.

# <img src="../doc/sources/Images/Machine_state_diagram.png" width="700">

# #### Instantiate the machine

# In[ ]:


pandora_machine = PandoraMachine()


# #### Define pipeline configuration

# In[ ]:


user_pipeline_cfg = {
    'pipeline':{
        "right_disp_map": {
            "method": "accurate"
        },
        "matching_cost" : {
            "matching_cost_method": "zncc",
            "window_size": 5,
            "subpix": 4
        },
        "aggregation": {
            "aggregation_method": "cbca"
        },
        "disparity": {
            "disparity_method": "wta",
            "invalid_disparity": "NaN"
        },
        "refinement": {
                "refinement_method": "quadratic"
        },
            "filter": {
                "filter_method": "median"
        },
            "validation": {
                "validation_method": "cross_checking"
        }
    }
}


# Disparity interval used

# In[ ]:


disp_min = -60
disp_max = 0


# #### Check the pipeline configuration and sequence of steps

# In[ ]:


checked_cfg = check_pipeline_section(user_pipeline_cfg, pandora_machine)


# In[ ]:


print(checked_cfg)


# In[ ]:


pipeline_cfg = checked_cfg['pipeline']


# #### Estimate the memory consumption of the pipeline

# In[ ]:


min_mem_consump, max_mem_consump = memory_consumption_estimation(user_pipeline_cfg, [img_left_path, disp_min, disp_max], pandora_machine)
print("Estimated maximum memory consumption between {:.2f} GiB and {:.2f} GiB".format(min_mem_consump, max_mem_consump))


# #### Prepare the machine 

# In[ ]:


pandora_machine.run_prepare(pipeline_cfg, img_left, img_right, disp_min, disp_max)


# #### Trigger the machine step by step

# In[ ]:


plot_state_machine(pandora_machine)


# Run matching cost 

# In[ ]:


pandora_machine.run('matching_cost', pipeline_cfg)


# In[ ]:


plot_state_machine(pandora_machine)


# In[ ]:


pandora_machine.run('aggregation', pipeline_cfg)


# In[ ]:


plot_state_machine(pandora_machine)


# Run disparity 

# In[ ]:


pandora_machine.run('disparity', pipeline_cfg)


# In[ ]:


plot_state_machine(pandora_machine)


# Run refinement 

# In[ ]:


pandora_machine.run('refinement', pipeline_cfg)


# In[ ]:


plot_state_machine(pandora_machine)


# Run filter

# In[ ]:


pandora_machine.run('filter', pipeline_cfg)


# In[ ]:


plot_state_machine(pandora_machine)


# Run validation

# In[ ]:


pandora_machine.run('validation', pipeline_cfg)


# In[ ]:


plot_state_machine(pandora_machine)


# Visualize output disparity map

# In[ ]:


plot_image(pandora_machine.left_disparity.disparity_map, "Left disparity map", output_dir, cmap=pandora_cmap())