#%%
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Build the hydrologic database used by the Generic PEM from a global DEM (.tif).

Runs "dephier.exe" to get the depression hierarchy, stitches the map east-west,
rebuilds the depression tree and writes "hydrologic_database.nc", then runs the
Fortran executable "database" which adds the lake volume/area/elevation curves
and writes the final database "hydrologic_database_final.nc".

The parameters are set in the PARAMETERS block below. See README.md for the
build of the executables and the Python environment.
"""

## @file database.py
#  @author Alexandre Gauvain
#  @date 03/02/2023

# =============================================================================
# PARAMETERS
# =============================================================================
RESULTS_FOLDER = "/scratchu/agauvain/trunk/PEM/toolbox/hydrology/results"
TIF_FILE = "/home/agauvain/work/MarsThroughTime/data/Mars_MGS_MOLA_DEM_mosaic_global_463m.tif"
DEPHIER_EXE = 'dephier.exe'
DATABASE_EXE = 'database'
R = 3397.2 #  Radius of Mars in PCM : 3397.2 Km | Radius of Earth in PCM :
RESHAPE_VALUE = 128
ADD_VALUE = 8

import os
import gc
import shutil
import psutil
import time
import datetime as dt
import subprocess
import netCDF4
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
import pandas as pd
pd.options.mode.chained_assignment = None
import math
import rasterio
import warnings
warnings.filterwarnings("ignore", category=rasterio.errors.NotGeoreferencedWarning)

import multiprocessing as mp
Ncpu = mp.cpu_count()

from pysheds.grid import Grid
#import whitebox
#wbt = whitebox.WhiteboxTools()
#wbt.set_compress_rasters(True)
#wbt.set_max_procs(Ncpu)
#wbt.verbose = True

class Database():
    def __init__(self, dem_path:str, results_path:str, R:float, dephier_exe:str, database_exe:str, reshape_value:int = 1, add_value:int = 4):
        """
        Build hydrological database from .tif of Mars DEM (MOLA).
        DEM: https://astrogeology.usgs.gov/search/map/Mars/GlobalSurveyor/MOLA/Mars_MGS_MOLA_DEM_mosaic_global_463m

        Parameters
        ----------
        dem_path : str
            Path of the DEM file .tif.
        dephier_exe : str
            Path of the depression hierarchy executable (default: <git>/bin/dephier.exe).
        database_exe : str
            Path of the Fortran database executable (default: <git>/src/hydro_database/database).
        """
        self.R = R #3389.5 Km Radius of Mars in Km
        self.start_time = time.time()
        print(dt.datetime.now().hour,':',dt.datetime.now().minute, ': Start of database construction')
        self.git_racine = os.path.dirname(os.path.dirname(os.path.dirname(os.path.join(os.getcwd(),__file__)))) 
        self.dem_path = dem_path
        self.dephier_exe = dephier_exe 
        self.database_exe = database_exe
        # Check inputs exist before running
        missing = []
        if not os.path.isdir(results_path):
            missing.append('results_path : ' + results_path)
        if not os.path.isfile(dem_path):
            missing.append('dem_path : ' + dem_path)
        for exe in [dephier_exe, database_exe]:
            if shutil.which(exe) is None:
                missing.append('executable (not found or not executable) : ' + exe)
        if missing:
            raise SystemExit('ERROR : missing inputs\n  - ' + '\n  - '.join(missing))
        self.reshape_value = reshape_value
        self.results_path = os.path.join(results_path,"Database_x"+str(self.reshape_value))
        self.tmp_path = os.path.join(results_path,"Database_x"+str(self.reshape_value),'tmp')
        self.figure_path = os.path.join(results_path,"Database_x"+str(self.reshape_value),"figure")
        if not os.path.exists(self.results_path):
            os.makedirs(self.results_path)
        if not os.path.exists(self.figure_path):
            os.makedirs(self.figure_path)
        if not os.path.exists(self.tmp_path):
            os.makedirs(self.tmp_path)
        self.add_value = add_value
        self.load_file()
        if reshape_value > 1: 
            self.reshape_dem(reshape_value,'min')
        
    def build_database(self):
        """
        Build hydrological database.
        """
        self.save_tif_file(self.elevation, 'elevation')
        elevation_sphere = np.concatenate((self.elevation[:,self.x_pixel-int(self.x_pixel/self.add_value):self.x_pixel],self.elevation), axis=1)
        elevation_sphere = np.concatenate((elevation_sphere,self.elevation[:,0:int(self.x_pixel/self.add_value)]), axis=1)
        self.save_tif_file(elevation_sphere,'elevation_dephier')
        del elevation_sphere
        self.d8_flow_direction()
        self.get_cells_area()
        #self.d8_flow_accumulation()
       	#self.flow_fill()
       	#self.display_flow_fill()
        self.depression_hierarchy()

    def cleanup(self):
        """
        Close the memory maps and delete the tmp folder.
        On NFS, a deleted file still opened by a memmap becomes a .nfsXXXX file
        and rmtree fails : the memmaps must be released before.
        """
        # list comprehension : no loop variable left holding a reference to a memmap
        for name in [n for n, v in vars(self).items() if isinstance(v, np.ndarray)]:
            delattr(self, name)
        gc.collect()
        shutil.rmtree(self.tmp_path, ignore_errors=True)
        if os.path.exists(self.tmp_path):
            print('WARNING : tmp folder not fully deleted (files still opened) : ' + self.tmp_path)

    def load_file(self):
        """
        Load .tif file : variables of data and georeferencing.
        """
        src = rasterio.open(self.dem_path)
        if self.reshape_value > 1:
            # Degraded DEM: no need to write the full-resolution DEM on disk, reshape_dem takes the min
            self.elevation = src.read(1)
        else:
            self.elevation = self.build_memory_map('elevation', np.int32(src.read(1)), 'int32')
        self.y_pixel = src.height
        self.x_pixel = src.width
        self.xmin = -180 #src.bounds.left
        self.xmax = 180 #src.bounds.right
        self.ymin = -90 #src.bounds.bottom
        self.ymax = 90 #src.bounds.top
        self.pixel_size = (self.xmax - self.xmin)/self.x_pixel #src.transform[0]
        self.profile = src.profile
        self.px_by_degree = int(self.elevation.shape[0]/180)
        self.latitude = self.build_memory_map('latitude', np.linspace(0,self.y_pixel,self.y_pixel) / -self.px_by_degree  + 90, 'float32') # pixels/degree
        self.longitude = self.build_memory_map('longitude', np.linspace(0,self.x_pixel-1,self.x_pixel) / self.px_by_degree , 'float32') # pixels/degree
        self.transform = rasterio.Affine(self.pixel_size, 0, self.xmin, 0, self.pixel_size * -1, self.ymax)
    
    def build_memory_map(self, name, data, dtype = 'float32'):
        memmap = np.memmap(os.path.join(self.tmp_path, name + '.dat'), dtype=dtype, mode='w+', shape=data.shape)
        memmap[:] = data[:]
        return memmap
    
    def reshape_dem(self, div_value:int, type_value:str='mean'):
        """
        Degrade the DEM by a div_value. The elevation data can be degraded by 
        the mean or the minimum of the group of considered pixels.

        Parameters
        ----------
        div_value : int
            Value to divide the number of pixel of the DEM.
        type_value : str, optional
            Type of degraded elevation data. The default is 'mean'.
        """
        # check if the array can be divisible by div_value and find closest
        # div_value if not
        if (self.x_pixel%div_value != 0) or (self.y_pixel%div_value != 0):
            div = False
            while div == False:
                div_value += 1
                if (self.x_pixel%div_value == 0) & (self.y_pixel%div_value == 0):
                    div = True
            self.reshape_value = div_value
            print('Warning : To reshape the DEM, the div_value need to be change to ' + str(div_value))
              
        # reshape the dem
        self.pixel_size = self.pixel_size * div_value
        self.y_pixel = int(self.y_pixel/div_value)
        self.x_pixel = int(self.x_pixel/div_value)
        self.latitude = np.linspace(0,self.y_pixel,self.y_pixel) / (-self.px_by_degree /div_value) + 90 # pixels/degree
        self.longitude = np.linspace(0,self.x_pixel-1,self.x_pixel) / (self.px_by_degree /div_value) # pixels/degree
        self.transform = rasterio.Affine(self.pixel_size, 0, self.xmin, 0, self.pixel_size * -1, self.ymax)
        if type_value == 'mean':
            self.elevation = np.int32(self.elevation).reshape(self.y_pixel, div_value, self.x_pixel, div_value).mean(axis=1).mean(axis=2)
        if type_value == 'min':
            self.elevation = np.int32(self.elevation.reshape(self.y_pixel, div_value, self.x_pixel, div_value).min(axis=(1, 3)))
        
    def get_cells_area(self):
        """
        Get the area of each cell of the  DEM
        Total surface of Mars is equal to 144 798 500 km²
        """
        # total_area = 144798500 * 1000000 # km² to m²
        lat_arr = np.array([self.latitude] * self.x_pixel).transpose()
        delta_long = np.abs(self.latitude[0] - self.latitude[1])
        delta_lat = np.abs(self.longitude[0] - self.longitude[1])
        self.area = self.build_memory_map('area', (np.pi * self.R * 1000 / 180)**2 * delta_long * delta_lat * np.cos(lat_arr * math.pi / 180), 'float32') # area of cells in m²
    
    def d8_flow_direction(self):
        dem_file = os.path.join(self.results_path,'elevation.tif')
        #output = os.path.join(self.results_path,'d8_flow_dir.tif')
        #wbt.d8_pointer(dem, output, esri_pntr=True)
        grid = Grid.from_raster(dem_file)
        dem = grid.read_raster(dem_file)
        #src = rasterio.open(os.path.join(self.results_path,'d8_flow_dir.tif'))
        fdir = grid.flowdir(dem)
        self.d8_flow_dir= self.build_memory_map('d8_flow_dir', fdir, fdir.dtype)
        #self.d8_flow_dir= self.d8_flow_dir[:,0+int(self.x_pixel/self.add_value):self.x_pixel+int(self.x_pixel/self.add_value)]
    
    def d8_flow_accumulation(self):
        dem = os.path.join(self.results_path,'elevation_dephier.tif')
        output = os.path.join(self.results_path,'d8_flow_acc.tif')
        wbt.d8_flow_accumulation(dem, output, out_type="cells",log=True, clip=1,)
        src = rasterio.open(output)
        self.d8_flow_acc = src.read(1)
        self.d8_flow_acc = self.d8_flow_acc[:,0+int(self.x_pixel/self.add_value):self.x_pixel+int(self.x_pixel/self.add_value)]
    
    def save_tif_file(self, data, name_file):
        """
        Save the .tif file in the results folder

        Parameters
        ----------
        data : array
            Data which is save in the first band of the .tif file.
        name_file : str
            Name of the file.
        """
        self.profile.update(
            width = data.shape[1],
            height = data.shape[0],
            transform = self.transform)
        with rasterio.open(os.path.join(self.results_path,name_file+'.tif'), 'w', **self.profile) as dst:
            dst.write(data, 1)
    
    def depression_hierarchy(self):
        """
        Barnes, R., Callaghan, K. L., & Wickert, A. D. (2020). Computing water 
        flow through complex landscapes–Part 2: Finding hierarchies in 
        depressions and morphological segmentations. Earth Surface Dynamics, 8(2), 431-445.
        source: https://esurf.copernicus.org/articles/8/431/2020/esurf-8-431-2020.html
        git: https://github.com/r-barnes/Barnes2019-DepressionHierarchy
        
        /!\/!\/!\ 
        The compiled file is in the bin folder but if recompiling is necessary :
        Before compiling, Add this part of code in the file Barnes2019-DepressionHierarchy/src/main.cpp 
        to extract all outputs depression hierarchy:
            
        {
        std::ofstream fgraph(out_name+"-graph-detail.csv");
        fgraph<<"child,parent,dep_vol,pit_cell,out_cell,out_elev,odep,geolink,oceanlink\n";
        for(unsigned int i=1;i<deps.size();i++){
        fgraph<<i<<","<<deps[i].parent<<","<<deps[i].dep_vol<<","<<deps[i].pit_cell<<","<<deps[i].out_cell<<","<<deps[i].out_elev<<","<<deps[i].odep<<","<<deps[i].geolink<<",";
        fgraph<<(deps[i].parent!=dh::NO_VALUE && (deps[i].parent==dh::OCEAN || !(deps[deps[i].parent].lchild==i || deps[deps[i].parent].rchild==i)))<<"\n";
        }
        }
        
        /!\/!\/!\
        """
        print('########## RUN DEPRESSION HIERARCHY C++ ##########')
        self.dephier_cmd = [self.dephier_exe,
                            os.path.join(self.results_path,'elevation_dephier.tif'),
                            os.path.join(self.results_path,'dephier'),
                            str(np.max(self.elevation))]
        print(self.dephier_cmd)
        self.run_command(' '.join(self.dephier_cmd), silent=False)
        print('########## LOAD DEPRESSION HIERARCHY FILE ##########')
        self.load_dephier_files()
        print('########## REBUILD DEPRESSION HIERARCHY ##########')
        self.rectify_and_resave_dephier_files()
        print('########## SAVE TMP HYDROLOGIC DATABASE ##########')
        self.save_hydrologic_database()
        print('########## RUN COMPLETE HYDROLOGIC DATABASE FORTRAN ##########')
        self.database_cmd = [self.database_exe,
                            os.path.join(self.results_path,'hydrologic_database.nc'),
                            os.path.join(self.results_path)]
        print(self.database_cmd)
        self.run_command(' '.join(self.database_cmd), silent=False)
        print('########## LOAD HYDROLOGIC DATABASE ##########')
        self.load_netcdf_file()
        #self.save_leaves()
        print('Hydrologic Database is created !')
    
    def load_dephier_files(self):
        """
        Load dephier-label.tif and dephier-graph.csv
        """
        src = rasterio.open(os.path.join(self.results_path,'dephier-label.tif'))
        dephier_label = src.read(1)
        self.dephier_label = self.build_memory_map('dephier_label', dephier_label, dephier_label.dtype)
        del dephier_label
        self.nleaf = np.max(self.dephier_label)
        self.graph = self.read_graph_csv(os.path.join(self.results_path,'dephier-graph-detail.csv'))
        print(self.graph)
        self.graph = self.build_memory_map('graph', self.graph, self.graph.dtype)
        self.graph_col = self.graph.dtype.names
        print('Number of leaf depressions and meta-depressions :',len(self.graph))
        self.dephier_inf_value = self.graph['odep'].max()
        self.print_memory_usage()
        #self.graph[self.graph == self.dephier_inf_value] = None
        #self.graph = self.graph.astype(int)

    def read_graph_csv(self, path):
        """
        Fast equivalent of np.genfromtxt(path, dtype=int, delimiter=',', names=True).
        Falls back on genfromtxt if a value is not an integer (genfromtxt gives -1 in this case).
        """
        try:
            df = pd.read_csv(path, dtype=np.int64)
        except (ValueError, OverflowError):
            return np.genfromtxt(path, dtype=int, delimiter=',', names=True)
        graph = np.empty(len(df), dtype=[(name, np.int64) for name in df.columns])
        for name in df.columns:
            graph[name] = df[name].to_numpy()
        return graph

        
    def load_netcdf_file(self):
        """
        Load hydrologic_database.nc
        """
        # Open the NetCDF file
        nc_file = netCDF4.Dataset(os.path.join(self.results_path,'hydrologic_database.nc'), 'r')
        # Access a variable and print its data
        self.dephier_graph['child'] = nc_file.variables['child'][:]
        self.dephier_graph['parent'] = nc_file.variables['parent'][:]
        self.dephier_graph['rchild'] = nc_file.variables['rchild'][:]
        self.dephier_graph['lchild'] = nc_file.variables['lchild'][:]
        self.dephier_graph['geolink'] = nc_file.variables['geolink'][:]
        self.dephier_graph['sibling'] = nc_file.variables['sibling'][:]
        self.dephier_graph['out_cell'] = nc_file.variables['out_cell'][:]
        self.dephier_graph['pit_cell'] = nc_file.variables['pit_cell'][:]
        self.dephier_graph['lake_lon'] = nc_file.variables['lake_lon'][:]
        self.dephier_graph['lake_lat'] = nc_file.variables['lake_lat'][:]
        self.dephier_graph['watershed_area'] = nc_file.variables['watershed_area'][:]
        self.dephier_graph['Vmax'] = nc_file.variables['Vmax'][:]
        self.dephier_graph['Vmin'] = nc_file.variables['Vmin'][:]
        self.dephier_graph['Zmax'] = nc_file.variables['Zmax'][:]
        self.dephier_graph['Zmin'] = nc_file.variables['Zmin'][:]
        self.dephier_graph['imin'] = nc_file.variables['imin'][:]
        self.dephier_graph['imax'] = nc_file.variables['imax'][:]
        self.dephier_graph['jmin'] = nc_file.variables['jmin'][:]
        self.dephier_graph['jmax'] = nc_file.variables['jmax'][:]
        
        self.lake_perc = nc_file.variables['lake_perc'][:]
        self.lake_vol = nc_file.variables['lake_vol'][:]
        self.lake_elev = nc_file.variables['lake_elev'][:]
        self.lake_area = nc_file.variables['lake_area'][:]
        
        # Close the file when done
        nc_file.close()
                
    def rectify_and_resave_dephier_files(self):
        """
        Rectify, check and save dephier-label.tif and dephier-graph-detail.csv
        (vectorised version of database_v8: same results)
        """
        pad = int(self.x_pixel/self.add_value)
        W = self.x_pixel
        H, Wall = self.dephier_label.shape
        print('(1) Rectify Label Map')
        self.start_time = time.time()
        self.label = self.build_memory_map('label',self.dephier_label[:,pad:W+pad], 'int')
        self.label_add_left = np.array(self.dephier_label[:,0:pad], dtype='int')
        self.label_add_right = np.array(self.dephier_label[:,W+pad::], dtype='int')
        self.label_add = np.concatenate((self.label_add_left, self.label_add_right), axis=1)
        self.label_conc_right = np.array(self.label[:,0:pad])
        self.label_conc_left = np.array(self.label[:,W-pad::])
        self.label_conc = np.concatenate((self.label_conc_left, self.label_conc_right), axis=1)
        self.print_memory_usage()
        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")

        """
        +-----------+--------------+--------------+--------------+-------------+
        |           |                                            |             |
        |           |                                            |             |
        |           |                label_map                   |             |
        |           |                                            |             |
        |           |                                            |             |
        |           |                                            |             |
        +-----------+--------------+--------------+--------------+-------------+
        <----------><------------->               <--------------><------------>
        label_add_left|label_conc_right         label_conc_right|label_add_right

        """
        # Rectify label number 0 on the map
        print('(2) Correct Label 0 on the map')
        self.start_time = time.time()
        idx = np.where(self.label == 0)
        for j in range(0, len(idx[0])):
            row = idx[0][j]
            col = idx[1][j]
            row_new, col_new = self.find_neighbor_cell(row, col, self.d8_flow_dir[row,col])
            self.label[row,col] = self.label[row_new, col_new]
        del idx
        self.print_memory_usage()
        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")

        #Rectify label map to connect west and east side
        print('(3) Find leafs connected')
        self.start_time = time.time()
        labels = np.array(list(set(zip(self.label_add_left[:,-1], self.label[:,self.x_pixel-1]))))
        idx_list = labels[:,0]
        idx_r_list = labels[:,1]
        del labels
        self.print_memory_usage()
        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")

        print('(4) Replace connected leafs on label map')
        self.start_time = time.time()
        # Same result as the sequential loop
        #   for i: label[:, 0:pad][label[:, 0:pad] == idx_list[i]] = idx_r_list[i]
        # but the successive replacements are applied on the label values (lookup table)
        # instead of scanning the map for each pair.
        region = self.label[:, 0:pad]
        members = {}   # current value -> original values which have this value
        touched = set()
        for a, b in zip(idx_list.tolist(), idx_r_list.tolist()):
            if a == b:
                continue
            if a in members:
                m = members.pop(a)
            elif a not in touched:
                touched.add(a)
                m = [a]
            else:
                continue
            if b in members:
                members[b].extend(m)
            elif b not in touched:
                touched.add(b)
                members[b] = [b] + m
            else:
                members[b] = m
        if len(members) > 0:
            vals, inv = np.unique(region, return_inverse=True)
            new_vals = vals.copy()
            pos = {v: k for k, v in enumerate(vals.tolist())}
            for v, origs in members.items():
                for o in origs:
                    k = pos.get(o)
                    if k is not None:
                        new_vals[k] = v
            region[...] = new_vals[inv].reshape(region.shape)
        del region
        self.print_memory_usage()
        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")

        # Pixel number : the dictionary {cell of the added bands: cell of the map} of v8
        # is replaced by an arithmetic function (see self.conc_cell)
        print('(5) Generate Id cells')
        self.start_time = time.time()
        self.Wall = Wall
        self.pad = pad
        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")

        print('(6) Find pairs of depressions')
        self.start_time = time.time()
        # Same array as np.array([[a, b, count] for (a, b), count in Counter(zip(label_add, label_conc)).items()])
        a = self.label_add.ravel()
        b = self.label_conc.ravel()
        M = int(b.max()) + 1
        key = a * M + b
        uniq, first, counts = np.unique(key, return_index=True, return_counts=True)
        order = np.argsort(first, kind='stable')   # order of first occurrence (Counter order)
        uniq = uniq[order]
        conc_arr = np.column_stack((uniq // M, uniq % M, counts[order])).astype('int')
        del a, b, key, uniq, first, counts, order
        conc_arr = conc_arr[conc_arr[:, 2].argsort()]
        keep = ~np.isin(conc_arr[:, 0], idx_r_list)
        conc_dict = dict(zip(conc_arr[keep, 0].tolist(), conc_arr[keep, 1].tolist()))
        self.print_memory_usage()
        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")

        print('(7) Clear graph by removing not connected leafs')
        self.start_time = time.time()
        mask = conc_arr[:, 0][~np.isin(conc_arr[:, 0], np.concatenate((idx_list, idx_r_list)))]
        self.graph = self.graph[~np.isin(self.graph['child'], mask)]
        del mask
        print("Execution time :", (time.time()-self.start_time)/60, "min")

        print('(8) Find spillover points')
        self.start_time = time.time()
        child_unique = len(np.unique(self.graph['child'])) == len(self.graph)
        # 'child' is unique -> same order as np.argsort(self.graph, order=['parent'])
        if child_unique:
            self.graph = self.graph[np.lexsort((self.graph['child'], self.graph['parent']))]
        else:
            self.graph = self.graph[np.argsort(self.graph, order=['parent'])]
        parents = np.asarray(self.graph['parent'])
        change = np.empty(len(parents), dtype=bool)
        change[0] = parents[0] != 0
        change[1:] = parents[1:] != parents[:-1]
        self.spillover_point = 1 + np.cumsum(change)
        # point_dict[k] of v8 == geolink_sorted[grp_start[k]:grp_start[k]+grp_len[k]]
        geolink_sorted = np.asarray(self.graph['geolink']).tolist()
        grp_len = np.bincount(self.spillover_point)
        grp_start = np.concatenate(([0], np.cumsum(grp_len)[:-1]))
        grp_len = grp_len.tolist()
        grp_start = grp_start.tolist()
        del parents, change
        print("Execution time :", (time.time()-self.start_time)/60, "min")

        # Build out_elev_list
        print('(9) Sort spillover points from lower to higher elevation')
        self.start_time = time.time()
        if child_unique:
            order = np.lexsort((self.graph['child'], self.graph['out_elev']))
        else:
            order = np.argsort(self.graph, order=['out_elev'])
        self.spillover_point = self.spillover_point[order]
        self.graph = self.graph[order]
        out_elev_list = self.graph['out_elev'] # List of overflow elevation
        out_cell_list = self.graph['out_cell'] # List of overflow ID cell
        out_point_list = self.spillover_point
        _, idx = np.unique(out_point_list[0:-1], return_index=True)
        out_cell = out_cell_list[np.sort(idx)]
        out_elev = out_elev_list[np.sort(idx)]
        out_point = out_point_list[np.sort(idx)]
        print("Execution time :", (time.time()-self.start_time)/60, "min")

        print('(10) Rectify depression volume, pit cell and pit_elevation')
        keys = ['child', 'parent','rchild', 'lchild', 'geolink', 'sibling', 'out_cell', 'i_out', 'j_out', 'pit_cell', 'lake_lon', 'lake_lat', 'watershed_area', 'Vmax', 'Vmin', 'Zmax', 'Zmin', 'imin', 'imax', 'jmin', 'jmax', 'check']
        keys_dtype = [int, int, int, int, int, int, int, int, int, int, float, float, float, float, float, float, float, int, int, int, int, int]
        np_dtype = []
        for i in range(0, len(keys)):
            if keys_dtype[i] == int:
                np_dtype.append((keys[i], np.int64))
            if keys_dtype[i] == float:
                np_dtype.append((keys[i], np.float64))

        # Rebuild graph table
        print('(11) Rebuild depression hierarchy graph table')
        self.start_time = time.time()
        # Outlet cells: padded map -> map (vectorised self.cell_conc_dict + self.find_out_cell)
        cells = self.conc_cell(np.asarray(out_cell, dtype=np.int64))
        col2, col1 = Wall, W
        x1 = cells - ((cells // col2) * (col2 - col1)) - ((col2 - col1) / 2) + 1
        i_out_arr = ((x1 - 1) // col1).astype(np.int64).tolist()
        j_out_arr = ((x1 - 1) % col1).astype(np.int64).tolist()
        new_cell_arr = x1.astype(np.int64).tolist()
        out_point = out_point.tolist()
        out_elev = out_elev.tolist()

        nleaf = int(self.nleaf)
        branch = nleaf + 1
        # Union-find with path compression: up[node] -> ancestor (up[root] == root).
        # find(x) == Tree.find_active_parent(x) of v8, without walking the whole branch each time.
        up = {}
        rec_child, rec_parent, rec_sibling, rec_geolink = [], [], [], []
        rec_cell, rec_zmax, rec_i, rec_j = [], [], [], []
        n_nodes = 0
        last_node = None

        def find(x):
            r = x
            p = up[r]
            while p != r:
                r = p
                p = up[r]
            while x != r:
                nx = up[x]
                up[x] = r
                x = nx
            return r

        for i in range(0, len(out_point)):
            point = out_point[i]
            if grp_len[point] != 2:
                continue
            s = grp_start[point]
            l0 = geolink_sorted[s]
            l1 = geolink_sorted[s + 1]
            l0 = conc_dict.get(l0, l0)
            l1 = conc_dict.get(l1, l1)
            c0 = find(l0) if l0 in up else l0
            c1 = find(l1) if l1 in up else l1
            if c0 != c1:
                elev = out_elev[i]
                cell = new_cell_arr[i]
                i_out = i_out_arr[i]
                j_out = j_out_arr[i]
                for l, c, oc, ol in ((l0, c0, c1, l1), (l1, c1, c0, l0)):
                    if c > nleaf:
                        node = c
                        if node not in up:
                            raise KeyError(node)
                    else:
                        node = l
                        up[node] = node
                        n_nodes += 1
                        last_node = node
                    up[node] = branch
                    rec_child.append(node); rec_parent.append(branch); rec_sibling.append(oc)
                    rec_geolink.append(ol); rec_zmax.append(elev); rec_cell.append(cell)
                    rec_i.append(i_out); rec_j.append(j_out)
                up[branch] = branch
                n_nodes += 1
                last_node = branch
                branch += 1
        del up
        # v8 keeps all the nodes except the last one added (the root): all of them must have a parent
        if n_nodes - len(rec_child) != 1:
            raise ValueError('Depression hierarchy: %d nodes without parent' % (n_nodes - len(rec_child)))

        self.dephier_graph = np.zeros(len(rec_child), dtype=np_dtype)
        self.dephier_graph['child'] = rec_child
        self.dephier_graph['parent'] = rec_parent
        self.dephier_graph['sibling'] = rec_sibling
        self.dephier_graph['geolink'] = rec_geolink
        self.dephier_graph['out_cell'] = rec_cell
        self.dephier_graph['Zmax'] = rec_zmax
        self.dephier_graph['i_out'] = rec_i
        self.dephier_graph['j_out'] = rec_j
        del rec_child, rec_parent, rec_sibling, rec_geolink, rec_cell, rec_zmax, rec_i, rec_j

        print('\r' +"Execution time :", (time.time()-self.start_time)/60, "min")
        # 'child' is unique -> same order as np.sort(self.dephier_graph, order='Zmax')
        self.dephier_graph = self.dephier_graph[np.lexsort((self.dephier_graph['child'], self.dephier_graph['Zmax']))]
        self.bigest_parent = self.dephier_graph['parent'].max()
        print('(12) Multiprocessing : Complete database')

        self.start_time = time.time()
        self.lake_dis = 11
        self.lake_perc= np.around(np.linspace(0,1,self.lake_dis),2)
        self.lake_vol = self.build_memory_map('lake_vol', np.zeros((len(self.dephier_graph), self.lake_dis), dtype='float64'), 'float64')
        self.lake_elev = self.build_memory_map('lake_elev', np.zeros((len(self.dephier_graph), self.lake_dis), dtype='float64'), 'float64')
        self.lake_area = self.build_memory_map('lake_area', np.ones((len(self.dephier_graph), self.lake_dis), dtype='float64'), 'float64')

    def conc_cell(self, cells):
        """
        Id of a cell of the padded map (elevation_dephier) -> Id of the same cell inside the map
        (same as self.cell_conc_dict of database_v8, cells outside the added bands are unchanged)
        """
        cells = np.asarray(cells, dtype=np.int64)
        H = self.dephier_label.shape[0]
        r, c = np.divmod(cells, self.Wall)
        valid = (cells >= 0) & (cells < H * self.Wall)
        left = valid & (c < self.pad)
        right = valid & (c >= self.pad + self.x_pixel)
        out = cells.copy()
        out[left] = cells[left] + self.x_pixel
        out[right] = cells[right] - self.x_pixel
        return out

    def find_out_cell(self, x2, col2, col1):
        x1 = x2 - ((x2 // col2) * (col2 - col1)) - ((col2 - col1) / 2) + 1
        i = (x1 - 1) // col1
        j = (x1 - 1) % col1
        return int(x1), int(i), int(j)
    
    def print_memory_usage(self):
        # Get memory usage statistics
        memory_stats = psutil.virtual_memory()
        # Print memory usage information
        # print(f"Total Memory: {memory_stats.total / (1024 ** 3):.2f} GB")
        # print(f"Available Memory: {memory_stats.available / (1024 ** 3):.2f} GB")
        # print(f"Used Memory: {memory_stats.used / (1024 ** 3):.2f} GB")
        print(f"Memory Usage Percentage: {memory_stats.percent:.2f}%")
    
    def array2string(self, array, delimiter=','):
        txt = delimiter.join([str(i) for i in array])
        return txt
   
    def smallestbox(a):
        r = a.any(1)
        if r.any():
            m,n = a.shape
            c = a.any(0)
            out = a[r.argmax():m-r[::-1].argmax(), c.argmax():n-c[::-1].argmax()]
        else:
            out = np.empty((0,0),dtype=bool)
        return out  
                    
    def find_leaf(self, parent):
        childs = []
        end = False
        child = list(self.dephier_graph['child'][self.dephier_graph['parent'] == parent])
        while end == False:
            if np.sum(child < self.nleaf) == len(child):
                end = True
            for c in range(0,len(child)):
                if child[0] <= self.nleaf:
                    childs.append(child[0])
                if child[0] > self.nleaf:
                   parent_child = self.dephier_graph['child'][self.dephier_graph['parent'] == child[0]]
                   for i in parent_child:
                       child.append(i)
                child.remove(child[0])
        return childs
    
    def find_leaves(self, parent):
        leaves_list = self.tree.leaves(parent)
        return [x.identifier for x in leaves_list]
        
    def check_dephier_graph(self):
        # Vectorised version of the v8 check (O(N log N) instead of O(N^2)), same messages
        child = np.asarray(self.dephier_graph['child'])
        ok = {}
        for key in ['parent', 'sibling', 'geolink']:
            ok[key] = np.isin(self.dephier_graph[key], child)
        if not ok['sibling'].all():
            # v8 stops with an IndexError when the sibling is not in the table
            raise IndexError('sibling ' + str(self.dephier_graph['sibling'][~ok['sibling']][0]) + ' not found')
        for k in np.where(~(ok['parent'] & ok['sibling'] & ok['geolink']))[0]:
            i = child[k]
            if not ok['parent'][k]:
                print('Check graph : an issue has be found for the child ' + str(i) + ' (parent = ' +str(self.dephier_graph['parent'][k:k+1])+ ')')
            if not ok['geolink'][k]:
                print('Check graph : an issue has be found for the child ' + str(i) + ' (geolink = ' +str(self.dephier_graph['geolink'][k:k+1])+ ')')
        print('Check done!')
    def check_dephier_outlet_cell(self):
        for l in self.dephier_graph.child.values:
            if l < self.nleaf:
                map_check = np.zeros((self.y_pixel, self.x_pixel))
                leaf = self.dephier_graph[self.dephier_graph['child'] == l]    
                parent = self.dephier_graph['parent'] [self.dephier_graph['child'] == l].values[0]
                
                nb = leaf.out_cell.values[0]
                map_check = np.zeros((self.y_pixel, self.x_pixel))
                
                if l > self.nleaf:
                    leafs = self.find_leaf(l)
                    for i in leafs:
                        map_check[self.label == i] = 1 
                if l < self.nleaf:
                  map_check[self.label == l] = 1
                try: 
                    brother = self.dephier_graph[(self.dephier_graph['parent']  == parent) & (self.dephier_graph['child'] != l)].child.values[0]
                    if brother > self.nleaf:
                        leafs = self.find_leaf(brother)
                        for i in leafs:
                            map_check[self.label == i] = 2 
                    if brother < self.nleaf:
                        map_check[self.label == brother] = 2
                except:pass
                map_check[self.cell_number == nb] = -1
                #self.display_dem(map_check, 'out_cell Id: '+str(l) + ' | ' + 'Connected depression: '+str(brother), cmap='jet', vmin=-1, vmax=2)
    
    def build_gif_dephier(self):
        import imageio
        display_list = self.dephier_graph.child[self.dephier_graph['child'] > self.nleaf].values[:-1]
        fig_created = []
        for i in display_list:
            parent = self.dephier_graph['parent'] [self.dephier_graph['child'] == i].values[0]
            brother = self.dephier_graph[(self.dephier_graph['parent']  == parent) & (self.dephier_graph['child'] != i)].child.values[0]
            if i>brother:
                fig_created.append(i)
        self.check_dephier_dep_vol(display=fig_created,save=True)
        # Build GIF
        with imageio.get_writer(os.path.join(self.figure_path,'dep.gif'), mode='I') as writer:
            for filename in fig_created:
                image = imageio.imread(os.path.join(self.figure_path,str(filename)+'.png'))
                writer.append_data(image)
                os.remove(os.path.join(self.figure_path,str(filename)+'.png'))
    
    def check_dephier_dep_vol(self, display='all',save=False):
        if display == 'all':
            display_list = self.dephier_graph['child']
        if display == 'leaf':
            display_list = self.dephier_graph['child'][self.dephier_graph['child'] <= self.nleaf]
        if display == 'parent':
            display_list = self.dephier_graph['child'][self.dephier_graph['child'] > self.nleaf]
        if display not in ['all','leaf','parent']:
            display_list=display
        
        map_check = np.zeros((self.y_pixel, self.x_pixel))
       
        for l in display_list:
            leaf = self.dephier_graph[self.dephier_graph['child'] == l]
            parent = self.dephier_graph['parent'][self.dephier_graph['child'] == l]
            brother = self.dephier_graph['child'][(self.dephier_graph['parent'] == parent) & (self.dephier_graph['child'] != l)]
            out_cell = leaf['out_cell']
            out_elev = leaf['Zmax']
            
            if l > self.nleaf:
                leafs = self.find_leaf(l)
                for i in leafs:
                    map_check[self.label == i] = 1 
                    map_check[(self.label == i) & (self.elevation <= out_elev)] = 2
            if l <= self.nleaf:
              map_check[self.label == l] = 1
              map_check[(self.label == l) & (self.elevation <= out_elev)] = 2
              
            
            if brother > self.nleaf:
                leafs = self.find_leaf(brother)
                for i in leafs:
                    map_check[self.label == i] = -2 
            if brother < self.nleaf:
                map_check[self.label == brother] = -2
                
            #map_check[self.cell_number == out_cell] = -1
            
            
            #self.display_dem(map_check, 'Id: '+str(l) + ' | ' + 'sibling: '+str(brother), cmap='jet_r', vmin=-2, vmax=2,save=save, save_name=str(l))
            map_check[map_check == -1] = 0
            map_check[map_check == -2] = 0
    
    def check_dephier_graph_label(self):
        """
        check to be sure that the rectify of label and graph is correct
        check the depression volume
        """
        leaves = self.dephier_graph['child'][self.dephier_graph['child'] <= self.nleaf]
        present = np.bincount(np.asarray(self.label).ravel()) > 0   # labels >= 0
        present = np.concatenate((present, np.zeros(max(0, int(leaves.max(initial=0)) + 1 - len(present)), dtype=bool)))
        for l in leaves[~present[leaves]]:
            print(l, "n'est pas sur la carte label")
        print('check done !')
    

    def find_neighbor_cell(self,row,col,pointer):
        if pointer == 1:
            row -= 1
            col += 1
        if pointer == 2:
            col += 1
        if pointer == 4:
            row += 1
            col += 1
        if pointer == 8:
            row += 1
        if pointer == 16:
            row += 1
            col -= 1
        if pointer == 32:
            col -= 1
        if pointer == 64:
            row -= 1
            col -= 1
        if pointer == 128:
            row -= 1
        if pointer == 0:
            row -= 1
            col += 1
        return row, col
    
    def build_dephier_tree(self, root = 0):
        """
        Build tree from treelib
        """
        from treelib import Tree
        self.tree = Tree()
        self.tree.create_node('Id='+str(root),root)
        for i in range(len(self.dephier_graph['child'])-1,-1, -1):
            self.tree.create_node(str(self.dephier_graph['child'][i]),self.dephier_graph['child'][i],parent=self.dephier_graph['parent'][i])
            
    def save_leaves(self):
        self.leaves = {}
        for i in range(0,len(self.dephier_graph['child'])):
            child = self.dephier_graph['child'][i]
            if child > self.nleaf:
                childs = []
                childs.append(self.dephier_graph['lchild'][i])
                childs.append(self.dephier_graph['rchild'][i])
                self.leaves[str(child)] = []
                for j in childs:
                    if j > self.nleaf:
                        for k in self.leaves[str(j)]:
                            self.leaves[str(child)].append(k)
                    if j <= self.nleaf:
                        self.leaves[str(child)].append(j)
                    
        import pickle
        with open(os.path.join(self.results_path,'leaves.dat'),'wb') as f:
            pickle.dump(self.leaves,f)
    
    def run_command(self, command, silent=False):
        process = subprocess.Popen(command, stdout=subprocess.PIPE, 
                           stderr=subprocess.PIPE, shell=True, text=True)
        prev_output = '00'
        while True:
            output = process.stdout.readline()
            if output == '' and process.poll() is not None:
                break
            if len(output.strip())>0:
                if not silent:
                    try:
                        if output.strip()[1] == '[' or output.strip()[-1] == '|':
                            print(output.strip(),end='\r')
                        else:
                            if prev_output[1] == '[' or prev_output[-1] == '|':
                                print('\n' + output.strip())
                            else:
                                print(output.strip())
                    except:pass
                prev_output = output.strip()
        rc = process.poll()
        return rc
    
    def remove_file(self, filename):
        try:
            os.remove(filename)
        except OSError:
            pass
                
    def save_hydrologic_database(self):
        self.remove_file(os.path.join(self.results_path,'hydrologic_database.nc'))
        ncfile = netCDF4.Dataset(os.path.join(self.results_path,'hydrologic_database.nc'), mode='w',format='NETCDF4')
        ncfile.setncattr('R', np.float64(self.R)) # Planet radius (km), read by the Fortran code
        ncfile.createDimension('lon', self.x_pixel)
        ncfile.createDimension('lat', self.y_pixel)
        ncfile.createDimension('depression', len(self.dephier_graph['child']))
        '''
        ncfile.createDimension('parent_depression', len(self.leaves_store.keys()))
        '''
        ncfile.createDimension('lake', self.lake_dis)
        '''
        ncfile.createDimension('leaf', self.len_leaves)
        '''
        lon_list = ncfile.createVariable('longitude', 'float64', ('lon'))
        lon_list[:] = self.longitude
        lat_list = ncfile.createVariable('latitude', 'float64', ('lat'))
        lat_list[:] = self.latitude
        area_list = ncfile.createVariable('area', 'float64', ('lat', 'lon'))
        area_list[:] = self.area
        '''
        cell_list = ncfile.createVariable('cell_number', 'float64', ('lat', 'lon'))
        cell_list[:] = self.cell_number
        '''
        label_list = ncfile.createVariable('label', 'float64', ('lat', 'lon'))
        label_list[:] = self.label
        elevation_list = ncfile.createVariable('elevation', 'float64', ('lat', 'lon'))
        elevation_list[:] = self.elevation
        flowdirs_list = ncfile.createVariable('flowdirs', 'int', ('lat', 'lon'))
        flowdirs_list[:] = self.d8_flow_dir
        '''
        leaves_len_list = ncfile.createVariable('leaves_len', 'int', ('parent_depression'))
        leaves_len_list[:] = self.leaves_len
        leaves_list = ncfile.createVariable('leaves', 'int', ('parent_depression', 'leaf'))
        leaves_list[:] = self.leaves
        '''
        lake_elev_list = ncfile.createVariable('lake_elev', 'float64', ('depression', 'lake'))
        lake_elev_list[:] = self.lake_elev
        lake_vol_list = ncfile.createVariable('lake_vol', 'float64', ('depression', 'lake'))
        lake_vol_list[:] = self.lake_vol
        lake_area_list = ncfile.createVariable('lake_area', 'float64', ('depression', 'lake'))
        lake_area_list[:] = self.lake_area
        lake_perc_list = ncfile.createVariable('lake_perc', 'float64', ('lake'))
        lake_perc_list[:] = self.lake_perc
        
        for i in self.dephier_graph.dtype.names:
            dtype = str(self.dephier_graph[i].dtype)
            if dtype == 'int64':
                netcdf_type = 'int'
            if dtype == 'float64':
                netcdf_type = 'float64'
            df_list = ncfile.createVariable(i, netcdf_type, ('depression'))
            df_list[:] = self.dephier_graph[i].astype(dtype)
        
        ncfile.close()
           
    def display_dem(self,data,text = 'Elevation (m)', cmap = None, vmin=None, vmax=None, save=False, save_name='fig'):
        fig = plt.figure(dpi=300, figsize=(5,5))
        ax = fig.add_subplot(aspect=True)
        long, lat = np.meshgrid(self.longitude, self.latitude)
        if vmin == None:
            vmin = np.min(data)
        if vmax == None:
            vmax = np.max(data)
        if cmap == None:
            cmap, vmin, vmax = self.get_cmap()
        data = ax.pcolormesh(long, lat, data, cmap=cmap, vmin=vmin, vmax=vmax)
        ax.contour(long, lat, self.elevation, 10, linewidths=0.5,
                   colors='black')
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="2.5%", pad=0.05)
        cbar = plt.colorbar(data, cax=cax)
        cbar.set_label(text)
        ax.set_xlabel('Longitude')
        ax.set_ylabel('Latitude')
        plt.tight_layout()
        if save is True:
            plt.savefig(os.path.join(self.figure_path, save_name+'.png'))
            fig.clear()
            plt.close(fig)
        else:
            plt.show()

    def get_cmap(self):
        color = pd.read_csv(os.path.join(self.git_racine,
                                         'molaTeam_contour_31rgb_steps.csv'))
        cdict = {}
        red = []
        green = []
        blue = []
        vmin = color['Start_elev'].min()
        vmax = color['Start_elev'].max()
        for i in range(0, len(color)):
            red.append([(color.iloc[i]['Start_elev'] - vmin)/(vmax-vmin),color.iloc[i]['Red']/255,color.iloc[i]['Red']/255])
            green.append([(color.iloc[i]['Start_elev'] - vmin)/(vmax-vmin),color.iloc[i]['Green']/255,color.iloc[i]['Green']/255])
            blue.append([(color.iloc[i]['Start_elev'] - vmin)/(vmax-vmin),color.iloc[i]['Blue']/255,color.iloc[i]['Blue']/255])
        cdict['red'] = red
        cdict['green'] = green
        cdict['blue'] = blue
        cmap = mcolors.LinearSegmentedColormap('mola', cdict)
        return cmap, vmin, vmax
    
class Tree():
    def __init__(self):
        self.Nodes = []
        self.Node_id = {}
        self.nNodes = 0
        
    def add_node(self, child, leaf = True):
        self.Nodes.append(Node(child,leaf))
        self.Node_id[child] = self.nNodes
        self.nNodes += 1

    def complete_node(self,child, parent, sibling, geolink, Zmax, cell, i_out, j_out):
        self.Nodes[self.Node_id[child]].parent = parent
        self.Nodes[self.Node_id[child]].sibling = sibling
        self.Nodes[self.Node_id[child]].geolink = geolink
        self.Nodes[self.Node_id[child]].Zmax = Zmax
        self.Nodes[self.Node_id[child]].cell = cell
        self.Nodes[self.Node_id[child]].i_out = i_out
        self.Nodes[self.Node_id[child]].j_out = j_out
        
    def find_active_parent(self, child):
        succes = False
        leaf = child
        while succes == False:
            Node = self.Nodes[self.Node_id[child]]
            if Node.parent == None:
                succes = True
            else:
                child = Node.parent
        return child
    
    def leaves_for_all_nodes(self):
        count = 0
        for n in self.Nodes:
            if n.leaf == True:
                leaf = n.child
                print(count)
                count += 1
                succes = False
                while succes == False:
                    if n.parent == None:
                        succes = True
                    else:
                        child = n.parent
                        n = self.Nodes[self.Node_id[child]]
                        n.leaves.append(leaf)
        
class Node():
    def __init__(self, child, leaf):
        self.child = child
        self.leaf = leaf
        self.parent = None
        self.lchild = None
        self.rchild = None
        self.sibling = None
        self.geolink = None
        self.Zmax = None
        self.cell = None
        self.i_out = None
        self.j_out = None
        self.leaves = []
        
    
    def update_params(self, parent, sibling, geolink, Zmax, cell, i_out, j_out):
        self.parent = parent
        self.sibling = sibling
        self.geolink = geolink
        self.Zmax = Zmax
        self.cell = cell
        self.i_out = i_out
        self.j_out = j_out
    
if __name__ == "__main__":
    DB = Database(TIF_FILE, RESULTS_FOLDER, R, dephier_exe = DEPHIER_EXE, database_exe = DATABASE_EXE, reshape_value = RESHAPE_VALUE, add_value = ADD_VALUE)
    DB.build_database()
    #DB.display_dem(DB.elevation)
    #DB.display_dem(DB.label, text = 'ID Label', cmap='jet')
    DB.check_dephier_graph_label()
    DB.check_dephier_graph()

    #DB.check_dephier_dep_vol(display='parent')
    DB.cleanup()