import os
from pathlib import Path
from abc import ABCMeta, abstractmethod
import json
from json import JSONEncoder
import numpy as np
from pyflowline.classes.vertex import pyvertex
from pyflowline.classes.edge import pyedge
from pyflowline.classes.flowline import pyflowline
from pyflowline.classes.confluence import pyconfluence
from pyflowline.formats.read_flowline import read_flowline_geojson
from pyflowline.formats.read_nhdplus_flowline_shapefile import read_nhdplus_flowline_shapefile_attribute
from pyflowline.formats.read_nhdplus_flowline_shapefile import extract_nhdplus_flowline_shapefile_by_attribute
from pyflowline.formats.read_nhdplus_flowline_shapefile import track_nhdplus_flowline
from pyflowline.formats.convert_flowline_to_geojson import convert_flowline_to_geojson
from pyflowline.formats.export_flowline import export_flowline_to_geojson
from pyflowline.formats.export_vertex import export_vertex_to_geojson
from pyflowline.algorithms.auxiliary.text_reader_string import text_reader_string
from pyflowline.algorithms.split.find_flowline_vertex import find_flowline_vertex
from pyflowline.algorithms.split.find_flowline_confluence import find_flowline_confluence
from pyflowline.algorithms.split.split_flowline import split_flowline
from pyflowline.algorithms.split.split_flowline_to_edge import split_flowline_to_edge
from pyflowline.algorithms.merge.merge_flowline import merge_flowline
from pyflowline.algorithms.direction.correct_flowline_direction import correct_flowline_direction
from pyflowline.algorithms.loop.remove_flowline_loop import remove_flowline_loop
from pyflowline.algorithms.simplification.remove_small_river import remove_small_river
from pyflowline.algorithms.simplification.remove_returning_flowline import remove_returning_flowline
from pyflowline.algorithms.simplification.remove_duplicate_flowline import remove_duplicate_flowline
from pyflowline.algorithms.index.define_stream_order import define_stream_order
from pyflowline.algorithms.index.define_stream_segment_index import define_stream_segment_index
from pyflowline.algorithms.intersect.intersect_flowline_with_mesh import intersect_flowline_with_mesh

class BasinClassEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.float32):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        if isinstance(obj, list):
            pass
        if isinstance(obj, pyvertex):
            return json.loads(obj.tojson()) #lVertexID
        if isinstance(obj, pyedge):
            return obj.lEdgeID
        if isinstance(obj, pyflowline):
            return obj.lFlowlineID
        if isinstance(obj, pyconfluence):
            return obj.dAngle_upstream

        return JSONEncoder.default(self, obj)

class pybasin(object):
    lBasinID =1
    sBasinID=''
    lCellID_outlet=-1
    iFlag_debug = 0
    iFlag_disconnected =0
    iFlag_dam=0
    dLongitude_outlet_degree = -9999.
    dLatitude_outlet_degree = -9999.
    dAccumulation_threshold= 100000.0
    dThreshold_small_river = 10000
    dLength_flowline_filtered = 0.0
    dLength_flowline_simplified = 0.0
    dLength_flowline_conceptual = 0.0

    dArea_of_difference=0.0
    dDistance_displace = 0.0
    sWorkspace_output_basin=''
    sFilename_flowline_raw=''
    sFilename_flowline_filter=''
    sFilename_flowline_filter_geojson=''
    sFilename_dam=''
    sFilename_flowline_topo=''
    #before intersect
    sFilename_flowline_simplified=''
    sFilename_flowline_segment_index_before_intersect=''
    sFilename_flowline_conceptual=''
    sFilename_flowline_edge=''
    sFilename_basin_info=''
    sFilename_flowline_simplified_info=''
    sFilename_flowline_conceptual_info=''
    sFilename_confluence_simplified_info=''
    sFilename_confluence_conceptual_info=''
    aFlowline_basin_filtered=None
    aFlowline_basin_simplified=None
    aFlowline_basin_conceptual=None
    pVertex_outlet=None
    aConfluence_basin_simplified= None
    aConfluence_basin_conceptual= None

    def __init__(self, aParameter):

        if 'lBasinID' in aParameter:
            self.lBasinID             = int(aParameter['lBasinID'])
        else:
            self.lBasinID   = 1

        if 'lCellID_outlet' in aParameter:
            self.lCellID_outlet             = int(aParameter['lCellID_outlet'])
        else:
            self.lCellID_outlet   = -1

        if 'iFlag_disconnected' in aParameter:
            self.iFlag_disconnected             = int(aParameter['iFlag_disconnected'])
        else:
            self.iFlag_disconnected   = 0

        if 'iFlag_dam' in aParameter:
            self.iFlag_dam             = int(aParameter['iFlag_dam'])
        else:
            self.iFlag_dam   = 0

        if 'iFlag_debug' in aParameter:
            self.iFlag_debug             = int(aParameter['iFlag_debug'])
        else:
            self.iFlag_debug   = 0

        if 'dLongitude_outlet_degree' in aParameter:
            self.dLongitude_outlet_degree             = float(aParameter['dLongitude_outlet_degree'])
        else:
            self.dLongitude_outlet_degree   = -9999.

        if 'dLatitude_outlet_degree' in aParameter:
            self.dLatitude_outlet_degree             = float(aParameter['dLatitude_outlet_degree'])
        else:
            self.dLatitude_outlet_degree   = -9999.

        if 'dThreshold_small_river' in aParameter:
            self.dThreshold_small_river             = float(aParameter['dThreshold_small_river'])
        else:
            self.dThreshold_small_river   = 10000.0

        if 'dAccumulation_threshold' in aParameter:
            self.dAccumulation_threshold             = float(aParameter['dAccumulation_threshold'])
        else:
            self.dAccumulation_threshold = 100000.0

        if 'sFilename_flowline_raw' in aParameter:
            self.sFilename_flowline_raw = aParameter['sFilename_flowline_raw']
        else:
            self.sFilename_flowline_raw   = ''

        if 'sFilename_flowline_filter' in aParameter:
            self.sFilename_flowline_filter = aParameter['sFilename_flowline_filter']
        else:
            self.sFilename_flowline_filter   = ''

        if 'sWorkspace_output_basin' in aParameter:
            self.sWorkspace_output_basin = aParameter['sWorkspace_output_basin']
        else:
            self.sWorkspace_output_basin   = '.'

        Path(self.sWorkspace_output_basin).mkdir(parents=True, exist_ok=True)

        self.sFilename_flowline_filter_geojson = os.path.join(str(self.sWorkspace_output_basin ), "flowline_filter.geojson"  )

        if 'sFilename_dam' in aParameter:
            self.sFilename_dam = aParameter['sFilename_dam']
        else:
            self.sFilename_dam   = ''

        if 'sFilename_flowline_topo' in aParameter:
            self.sFilename_flowline_topo = aParameter['sFilename_flowline_topo']
        else:
            self.sFilename_flowline_topo   =''

        self.sBasinID = sBasinID = "{:03d}".format(self.lBasinID)

        #geojson
        self.sFilename_flowline_segment_index_before_intersect = 'flowline_segment_index_before_intersect.geojson'
        self.sFilename_flowline_simplified = 'flowline_simplified.geojson'
        self.sFilename_flowline_intersect  = 'flowline_intersect_mesh.geojson'
        self.sFilename_flowline_conceptual = 'flowline_conceptual.geojson'
        self.sFilename_flowline_edge = 'flowline_edge.geojson'
        self.sFilename_area_of_difference = 'area_of_difference.geojson'

        self.sFilename_basin_info = 'basin_info.json'
        self.sFilename_flowline_conceptual_info = 'flowline_conceptual_info.json'
        self.sFilename_flowline_simplified_info = 'flowline_simplified_info.json'
        self.sFilename_confluence_conceptual_info = 'confluence_conceptual_info.json'
        self.sFilename_confluence_simplified_info = 'confluence_simplified_info.json'
        return

    def flowline_simplification(self):

        sFilename_flowline_filter = self.sFilename_flowline_filter
        sFilename_flowline_filter_geojson = self.sFilename_flowline_filter_geojson
        aFlowline_basin_filtered, pSpatial_reference = read_flowline_geojson( sFilename_flowline_filter_geojson )
        sWorkspace_output_basin = self.sWorkspace_output_basin
        if self.iFlag_dam ==1:
            sFilename_dam = self.sFilename_dam
            aData_dam = text_reader_string(sFilename_dam, iSkipline_in =1,cDelimiter_in=',' )
            sFilename_flowline_topo = self.sFilename_flowline_topo
            aData_flowline_topo = text_reader_string(sFilename_flowline_topo, iSkipline_in =1,cDelimiter_in=',' )
            aFromFlowline = aData_flowline_topo[:,1].astype(int).ravel()
            aToFlowline = aData_flowline_topo[:,2].astype(int).ravel()
            sFilename_flowline_raw = self.sFilename_flowline_raw
            aNHDPlusID_filter = read_nhdplus_flowline_shapefile_attribute(sFilename_flowline_filter)
            aNHDPlusID_raw = read_nhdplus_flowline_shapefile_attribute(sFilename_flowline_raw)
            ndam = len(aData_dam)
            aNHDPlusID_dams_headwater = list()
            aNHDPlusID_dams_nonheadwater = list()
            for j in range(0, ndam):
                dLon = float(aData_dam[j][1])
                dLat = float(aData_dam[j][0])
                sDam = aData_dam[j][4]
                lNHDPlusID = int(aData_dam[j][5])
                aNHDPlusID_dams_headwater.append(lNHDPlusID)
                if lNHDPlusID in aNHDPlusID_filter:
                    #remove by id
                    for k in range(len(aFlowline_basin_filtered)):
                        if aFlowline_basin_filtered[k].lNHDPlusID == lNHDPlusID:
                            aFlowline_basin_filtered.pop(k)
                            break
                    pass
                else:
                    aNHDPlusID_dam_nonheadwater = track_nhdplus_flowline(aNHDPlusID_filter, aFromFlowline, aToFlowline, lNHDPlusID)
                    aNHDPlusID_filter = aNHDPlusID_filter + aNHDPlusID_dams_headwater+ aNHDPlusID_dam_nonheadwater
                    aNHDPlusID_dams_nonheadwater = aNHDPlusID_dams_nonheadwater + aNHDPlusID_dam_nonheadwater
            aFlowline_dams_headwater = extract_nhdplus_flowline_shapefile_by_attribute(sFilename_flowline_raw, aNHDPlusID_dams_headwater )
            for i in range(len(aFlowline_dams_headwater)):
                aFlowline_dams_headwater[i].iFlag_dam = 1
            aFlowline_dams_nonheadwater = extract_nhdplus_flowline_shapefile_by_attribute(sFilename_flowline_raw, aNHDPlusID_dams_nonheadwater )
            aFlowline_basin_filtered = aFlowline_basin_filtered + aFlowline_dams_headwater + aFlowline_dams_nonheadwater
        else:
            pass
        if self.iFlag_disconnected == 1:
            #not used anymore
            #aThreshold = np.full(2, 300.0, dtype=float)
            #aFlowline_basin_filtered = connect_disconnect_flowline(aFlowline_basin_filtered, aVertex, aThreshold)
            #sFilename_out = 'flowline_connect.geojson'
            #sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            #export_flowline_to_geojson(iFlag_projected, aFlowline_basin_filtered,pSpatial_reference_gcs, sFilename_out)
            pass
        else:
            pass

        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_before_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            self.export_flowline(aFlowline_basin_filtered, sFilename_out)
        #calculate length
        self.aFlowline_basin_filtered = aFlowline_basin_filtered
        self.dLength_flowline_filtered = self.calculate_flowline_length(aFlowline_basin_filtered)

        #simplification started
        aVertex = find_flowline_vertex(aFlowline_basin_filtered)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_vertex_without_confluence_before_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_vertex_to_geojson( aVertex, sFilename_out)
        aFlowline_basin_simplified = split_flowline(aFlowline_basin_filtered, aVertex)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_split_by_point_before_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson(aFlowline_basin_simplified, sFilename_out)
        #ues location to find outlet
        point= dict()
        point['dLongitude_degree'] = self.dLongitude_outlet_degree
        point['dLatitude_degree'] = self.dLatitude_outlet_degree
        pVertex_outlet=pyvertex(point)
        aFlowline_basin_simplified = correct_flowline_direction(aFlowline_basin_simplified,  pVertex_outlet )
        pVertex_outlet = aFlowline_basin_simplified[0].pVertex_end
        self.pVertex_outlet = pVertex_outlet
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_direction_before_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson( aFlowline_basin_simplified,  sFilename_out)
        #step 4: remove loops
        aFlowline_basin_simplified = remove_flowline_loop(aFlowline_basin_simplified)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_loop_before_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson( aFlowline_basin_simplified, sFilename_out)
        #using loop to remove small river, here we use 5 steps
        for i in range(3):
            sStep = "{:02d}".format(i+1)
            aFlowline_basin_simplified = remove_small_river(aFlowline_basin_simplified, self.dThreshold_small_river)
            if self.iFlag_debug ==1:
                sFilename_out = 'flowline_large_'+ sStep +'_before_intersect.geojson'
                sFilename_out =os.path.join(sWorkspace_output_basin, sFilename_out)
                export_flowline_to_geojson( aFlowline_basin_simplified,  sFilename_out)
            aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity = find_flowline_confluence(aFlowline_basin_simplified,  pVertex_outlet)
            if self.iFlag_debug ==1:
                sFilename_out = 'flowline_vertex_with_confluence_'+ sStep +'_before_intersect.geojson'
                sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
                export_vertex_to_geojson( aVertex,  sFilename_out, aAttribute_data=aConnectivity)
            aFlowline_basin_simplified = merge_flowline( aFlowline_basin_simplified,aVertex, pVertex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence  )
            if self.iFlag_debug ==1:
                sFilename_out = 'flowline_merge_'+ sStep +'_before_intersect.geojson'
                sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
                export_flowline_to_geojson( aFlowline_basin_simplified,  sFilename_out)
            if len(aFlowline_basin_simplified) == 1:
                break

        #the final vertex info
        aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity = find_flowline_confluence(aFlowline_basin_simplified,  pVertex_outlet)
        sFilename_out = 'vertex_simplified.geojson'
        sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
        export_vertex_to_geojson( aVertex,  sFilename_out, aAttribute_data=aConnectivity)

        #self.dLength_flowline_simplified = self.calculate_flowline_length(aFlowline_basin_simplified)
        aVertex = np.array(aVertex)
        aIndex_confluence = np.array(aIndex_confluence)
        if aIndex_confluence.size > 0:
            aVertex_confluence = aVertex[aIndex_confluence]
            self.aConfluence_basin_simplified = self.build_confluence(aFlowline_basin_simplified, aVertex_confluence)

        #build segment index
        aFlowline_basin_simplified, aStream_segment = define_stream_segment_index(aFlowline_basin_simplified)
        if self.iFlag_debug ==1:
            sFilename_out = self.sFilename_flowline_segment_index_before_intersect
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson(  aFlowline_basin_simplified, sFilename_out, \
                aAttribute_data=[aStream_segment], aAttribute_field=['iseg'], aAttribute_dtype=['int'])
        #build stream order
        aFlowline_basin_simplified, aStream_order = define_stream_order(aFlowline_basin_simplified)
        sFilename_out = self.sFilename_flowline_simplified
        sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
        export_flowline_to_geojson(  aFlowline_basin_simplified, sFilename_out, \
                aAttribute_data=[aStream_segment, aStream_order], aAttribute_field=['iseg','iord'], aAttribute_dtype=['int','int'])

        self.aFlowline_basin_simplified= aFlowline_basin_simplified
        return aFlowline_basin_simplified

    def reconstruct_topological_relationship(self, iMesh_type, sFilename_mesh):

        sWorkspace_output_basin = self.sWorkspace_output_basin
        sFilename_flowline = self.sFilename_flowline_simplified
        sFilename_flowline_in = os.path.join(sWorkspace_output_basin, sFilename_flowline)
        aFlowline_basin_simplified, pSpatial_reference = read_flowline_geojson( sFilename_flowline_in )

        sFilename_flowline_intersect = self.sFilename_flowline_intersect
        sFilename_flowline_intersect_out = os.path.join(sWorkspace_output_basin, sFilename_flowline_intersect)
        aCell, aCell_intersect_basin, aFlowline_intersect_all = intersect_flowline_with_mesh(iMesh_type, sFilename_mesh, \
            sFilename_flowline_in, sFilename_flowline_intersect_out)
        sFilename_flowline_filter_geojson = self.sFilename_flowline_filter
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_intersect_flowline_with_mesh.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson(aFlowline_intersect_all,  sFilename_out)

        point= dict()
        point['dLongitude_degree'] = self.dLongitude_outlet_degree
        point['dLatitude_degree'] = self.dLatitude_outlet_degree
        pVertex_outlet_initial=pyvertex(point)

        #from this point, aFlowline_basin is conceptual
        #segment based
        aFlowline_basin_conceptual, lCellID_outlet, pVertex_outlet = remove_returning_flowline(iMesh_type, aCell_intersect_basin, pVertex_outlet_initial)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_simplified_after_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson(aFlowline_basin_conceptual,  sFilename_out)

        #edge based
        aFlowline_basin_conceptual, aEdge = split_flowline_to_edge(aFlowline_basin_conceptual)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_edge_split_flowline_to_edge.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson( aFlowline_basin_conceptual,  sFilename_out)
        aFlowline_basin_conceptual = remove_duplicate_flowline(aFlowline_basin_conceptual)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_edge_remove_duplicate_flowline.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson( aFlowline_basin_conceptual,  sFilename_out)
        aFlowline_basin_conceptual = correct_flowline_direction(aFlowline_basin_conceptual,  pVertex_outlet )
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_edge_correct_flowline_direction.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson( aFlowline_basin_conceptual,  sFilename_out)
        aFlowline_basin_conceptual = remove_flowline_loop(  aFlowline_basin_conceptual )
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_edge_remove_flowline_loop.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_flowline_to_geojson( aFlowline_basin_conceptual,  sFilename_out)

        aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity\
            = find_flowline_confluence(aFlowline_basin_conceptual,  pVertex_outlet)
        if self.iFlag_debug ==1:
            sFilename_out = 'flowline_vertex_with_confluence_after_intersect.geojson'
            sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
            export_vertex_to_geojson( aVertex,  sFilename_out, aAttribute_data=aConnectivity)

        #segment based
        aFlowline_basin_conceptual = merge_flowline( aFlowline_basin_conceptual,aVertex, pVertex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence  )
        aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity\
            = find_flowline_confluence(aFlowline_basin_conceptual,  pVertex_outlet)
        aFlowline_basin_conceptual, aStream_segment = define_stream_segment_index(aFlowline_basin_conceptual)
        aFlowline_basin_conceptual, aStream_order = define_stream_order(aFlowline_basin_conceptual)

        #save confluence
        aVertex = np.array(aVertex)
        aIndex_confluence = np.array(aIndex_confluence)
        if aIndex_confluence.size > 0:
            aVertex_confluence = aVertex[aIndex_confluence]
            self.aConfluence_basin_conceptual = self.build_confluence(aFlowline_basin_conceptual, aVertex_confluence)

        #edge based
        aFlowline_basin_edge, aEdge = split_flowline_to_edge(aFlowline_basin_conceptual)
        sFilename_out = self.sFilename_flowline_edge
        sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
        export_flowline_to_geojson(  aFlowline_basin_edge, sFilename_out)

        sFilename_out = self.sFilename_flowline_conceptual
        sFilename_out = os.path.join(sWorkspace_output_basin, sFilename_out)
        export_flowline_to_geojson(  aFlowline_basin_conceptual, sFilename_out, \
            aAttribute_data=[aStream_segment, aStream_order], aAttribute_field=['iseg','iord'], aAttribute_dtype=['int','int'])

        self.aFlowline_basin_conceptual = aFlowline_basin_conceptual

        self.lCellID_outlet = lCellID_outlet
        self.dLongitude_outlet_degree = pVertex_outlet.dLongitude_degree
        self.dLatitude_outlet_degree = pVertex_outlet.dLatitude_degree

        return aCell_intersect_basin

    def build_confluence(self, aFlowline_basin_in, aVertex_confluence_in):
        #this can only be calculated for confluence
        aConfluence_basin=list()
        for pVertex in aVertex_confluence_in:
            aFlowline_upstream =list()
            for pFlowline in aFlowline_basin_in:
                pVertex_start = pFlowline.pVertex_start
                pVertex_end = pFlowline.pVertex_end
                if pVertex_end == pVertex:
                    aFlowline_upstream.append(pFlowline)
                    pass
                if pVertex_start == pVertex:
                    pFlowline_downstream=pFlowline

            pConfluence = pyconfluence(pVertex, aFlowline_upstream, pFlowline_downstream)
            aConfluence_basin.append(pConfluence)
        return aConfluence_basin

    def analyze(self):
        if self.aFlowline_basin_filtered is None:
            sFilename_flowline_filter = self.sFilename_flowline_filter
            sFilename_flowline_filter_geojson = self.sFilename_flowline_filter_geojson
            self.aFlowline_basin_filtered, pSpatial_reference = read_flowline_geojson( sFilename_flowline_filter_geojson )
            self.dLength_flowline_filtered = self.calculate_flowline_length(self.aFlowline_basin_filtered)

        point= dict()
        point['dLongitude_degree'] = self.dLongitude_outlet_degree
        point['dLatitude_degree'] = self.dLatitude_outlet_degree
        pVertex_outlet_initial=pyvertex(point)
        if self.aFlowline_basin_simplified is None:
            sFilename_flowline = self.sFilename_flowline_simplified
            sFilename_flowline_in = os.path.join(self.sWorkspace_output_basin, sFilename_flowline)
            aFlowline_simplified,pSpatial_reference = read_flowline_geojson( sFilename_flowline_in )
            read_flowline_geojson
            self.aFlowline_basin_simplified = aFlowline_simplified
            aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity\
            = find_flowline_confluence(self.aFlowline_basin_simplified,  pVertex_outlet_initial)
            aVertex = np.array(aVertex)
            aIndex_confluence = np.array(aIndex_confluence)
            if aIndex_confluence.size > 0:
                aVertex_confluence = aVertex[aIndex_confluence]
                self.aConfluence_basin_simplified = self.build_confluence(self.aFlowline_basin_simplified, aVertex_confluence)

        self.dLength_flowline_simplified = self.calculate_flowline_length(self.aFlowline_basin_simplified)

        if self.aFlowline_basin_conceptual is None:
            sFilename_flowline = self.sFilename_flowline_conceptual
            sFilename_flowline_in = os.path.join(self.sWorkspace_output_basin, sFilename_flowline)
            aFlowline_conceptual, pSpatial_reference = read_flowline_geojson( sFilename_flowline_in )
            self.aFlowline_basin_conceptual = aFlowline_conceptual

            aVertex, lIndex_outlet, aIndex_headwater,aIndex_middle, aIndex_confluence, aConnectivity\
            = find_flowline_confluence(self.aFlowline_basin_conceptual,  pVertex_outlet_initial)
            aVertex = np.array(aVertex)
            aIndex_confluence = np.array(aIndex_confluence)
            if aIndex_confluence.size > 0:
                aVertex_confluence = aVertex[aIndex_confluence]
                self.aConfluence_basin_conceptual = self.build_confluence(self.aFlowline_basin_conceptual, aVertex_confluence)

        self.dLength_flowline_conceptual = self.calculate_flowline_length(self.aFlowline_basin_conceptual)
        self.calculate_river_sinuosity()
        self.calculate_confluence_branching_angle()
        return

    def export(self):
        self.export_basin_info_to_json()
        self.export_flowline_info_to_json()
        self.export_confluence_info_to_json()
        return

    def export_flowline(self, aFlowline_in, sFilename_json_in,iFlag_projected_in = None,  pSpatial_reference_in = None):
        export_flowline_to_geojson(aFlowline_in, sFilename_json_in,\
            iFlag_projected_in= iFlag_projected_in, \
            pSpatial_reference_in = pSpatial_reference_in)

    def export_basin_info_to_json(self):
        sFilename_json = self.sFilename_basin_info
        sFilename_json = os.path.join(str(Path(self.sWorkspace_output_basin)  ) , sFilename_json  )

        aSkip = ['aFlowline_basin_filtered', \
                'aFlowline_basin_simplified','aFlowline_basin_conceptual','aConfluence_basin_simplified',
                'aConfluence_basin_conceptual']
        obj = self.__dict__.copy()
        for sKey in aSkip:
            obj.pop(sKey, None)
        with open(sFilename_json, 'w', encoding='utf-8') as f:
            sJson = json.dumps(obj, default=lambda o: o.__dict__,\
            sort_keys=True, \
                indent = 4, \
                    ensure_ascii=True, \
                        cls=BasinClassEncoder)
            f.write(sJson)
            f.close()
        return

    def export_flowline_info_to_json(self):
        iFlag_export_simplified=0
        if iFlag_export_simplified==1:
            sFilename_json = self.sFilename_flowline_simplified_info
            sFilename_json = os.path.join(str(Path(self.sWorkspace_output_basin)  ) , sFilename_json  )
            with open(sFilename_json, 'w', encoding='utf-8') as f:
                sJson = json.dumps([json.loads(ob.tojson()) for ob in self.aFlowline_basin_simplified], indent = 4)
                f.write(sJson)
                f.close()

        sFilename_json = self.sFilename_flowline_conceptual_info
        sFilename_json = os.path.join(str(Path(self.sWorkspace_output_basin)  ) , sFilename_json  )

        with open(sFilename_json, 'w', encoding='utf-8') as f:
            sJson = json.dumps([json.loads(ob.tojson()) for ob in self.aFlowline_basin_conceptual], indent = 4)
            f.write(sJson)
            f.close()
        return

    def export_confluence_info_to_json(self):
        iFlag_export_confluence =0
        if iFlag_export_confluence==1:
            sFilename_json = self.sFilename_confluence_simplified_info
            sFilename_json = os.path.join(str(Path(self.sWorkspace_output_basin)  ) , sFilename_json  )
            with open(sFilename_json, 'w', encoding='utf-8') as f:
                sJson = json.dumps([json.loads(ob.tojson()) for ob in self.aConfluence_basin_simplified], indent = 4)
                f.write(sJson)
                f.close()

        sFilename_json = self.sFilename_confluence_conceptual_info
        sFilename_json = os.path.join(str(Path(self.sWorkspace_output_basin)  ) , sFilename_json  )

        with open(sFilename_json, 'w', encoding='utf-8') as f:
            sJson = json.dumps([json.loads(ob.tojson()) for ob in self.aConfluence_basin_conceptual], indent = 4)
            f.write(sJson)
            f.close()
        return

    def tojson(self):
        aSkip = ['aFlowline_basin_filtered', \
                'aFlowline_basin_simplified','aFlowline_basin_conceptual','aConfluence_basin_simplified',
                'aConfluence_basin_conceptual']

        obj = self.__dict__.copy()
        for sKey in aSkip:
            obj.pop(sKey, None)

        sJson = json.dumps(obj, \
            sort_keys=True, \
                indent = 4, \
                    ensure_ascii=True, \
                        cls=BasinClassEncoder)
        return sJson

    def export_config_to_json(self, sFilename_output_in = None):
        #single basin
        if sFilename_output_in is not None:
            sFilename_output = sFilename_output_in
        else:
            sFilename_output = os.path.join(self.sWorkspace_output_basin, 'configuration_basin.json' )

        aSkip = ['aFlowline_basin_filtered', \
                'aFlowline_basin_simplified','aFlowline_basin_conceptual','aConfluence_basin_simplified',
                'aConfluence_basin_conceptual']
        obj = self.__dict__.copy()
        for sKey in aSkip:
            obj.pop(sKey, None)
        with open(sFilename_output, 'w', encoding='utf-8') as f:
            json.dump(obj, f,sort_keys=True, \
                ensure_ascii=False, \
                indent=4, \
                cls=BasinClassEncoder)
        return

    def convert_flowline_to_geojson(self):
        sFilename_raw = self.sFilename_flowline_filter
        sFilename_out = self.sFilename_flowline_filter_geojson
        print('This is the filtered flowline:', sFilename_raw )
        convert_flowline_to_geojson(1, sFilename_raw, sFilename_out)

    def calculate_flowline_length(self, aFlowline_in):
        dLength = 0.0
        nflowline = len(aFlowline_in)
        for i in range(nflowline):
            pFlowline= aFlowline_in[i]
            pFlowline.calculate_length()
            dLength = dLength + pFlowline.dLength
        return dLength

    def calculate_river_sinuosity(self):
        for pFlowline in self.aFlowline_basin_simplified:
            pFlowline.calculate_flowline_sinuosity()

        for pFlowline in self.aFlowline_basin_conceptual:
            pFlowline.calculate_flowline_sinuosity()

        return

    def calculate_confluence_branching_angle(self):
        for pConfluence in self.aConfluence_basin_simplified:
            pConfluence.calculate_branching_angle()
        for pConfluence in self.aConfluence_basin_conceptual:
            pConfluence.calculate_branching_angle()
        return


"""Classes representing different kinds of astronomical position."""

from numpy import array, arccos, clip, einsum, exp

from .constants import RAD2DEG, tau
from .data.spice import inertial_frames
from .functions import dots, from_polar, length_of, to_polar, rot_z
from .earthlib import compute_limb_angle, refract
from .relativity import add_aberration, add_deflection
from .timelib import Time
from .units import Distance, Velocity, Angle, _interpret_angle

_ECLIPJ2000 = inertial_frames['ECLIPJ2000']
_GALACTIC = inertial_frames['GALACTIC']

def build_position(position_au, velocity_au_per_d=None, t=None,
                   center=None, target=None, observer_data=None):
    if center == 0:
        cls = Barycentric
    elif center == 399:
        cls = Geocentric
    elif observer_data is not None:
        cls = Geometric
    else:
        cls = ICRF
    return cls(position_au, velocity_au_per_d, t, center, target, observer_data)

class ICRF(object):
    """An (x, y, z) position and velocity oriented to the ICRF axes.

    The ICRF is a permanent coordinate system that has superseded the
    old series of equinox-based systems like B1900 and B1950.  Its axes
    are aligned with the axes of J2000 to within 0.02 arcseconds, which
    is tighter than the accuracy of J2000 itself.

    """
    def __init__(self, position_au, velocity_au_per_d=None, t=None,
                 center=None, target=None, observer_data=None):
        self.t = t
        self.position = Distance(position_au)
        if velocity_au_per_d is None:
            self.velocity = None
        else:
            self.velocity = Velocity(velocity_au_per_d)
        # TODO: are center and target useful? Then why are they not
        # propagated down to Astrometric and Apparent positions?
        self.center = center
        self.target = target
        self.observer_data = observer_data

    def __repr__(self):
        return '<{0} position{1}{2}{3}{4}>'.format(
            self.__class__.__name__,
            '' if (self.velocity is None) else ' and velocity',
            '' if self.t is None else ' at date t',
            '' if self.center is None else ' center={0}'.format(self.center),
            '' if self.target is None else ' target={0}'.format(self.target),
        )

    def __sub__(self, body):
        """Subtract two ICRF vectors to produce a third."""
        # TODO: set center and target of result
        p = self.position.au - body.position.au
        if self.velocity is None or body.velocity is None:
            v = None
        else:
            v = body.velocity.au_per_d - self.velocity.au_per_d
        return ICRF(p, v, self.t)

    def distance(self):
        """Compute the distance from the origin to this position.

        >>> v = ICRF([1, 1, 0])
        >>> print(v.distance())
        1.41421 au

        """
        return Distance(length_of(self.position.au))

    def speed(self):
        """Compute the magnitude of the velocity vector.

        >>> v = ICRF([0, 0, 0], [1, 2, 3])
        >>> print(v.speed())
        3.74166 au/day

        """
        return Velocity(length_of(self.velocity.au_per_d))

    def radec(self, epoch=None):
        r"""Compute equatorial (RA, declination, distance)

        When called without a parameter, this returns standard ICRF
        right ascension and declination:

        >>> ra, dec, distance = ICRF([1, 2, 3]).radec()
        >>> print(ra, dec, distance, sep='\n')
        04h 13m 44.39s
        +53deg 18' 02.8"
        3.74166 au

        If you instead want the coordinates referenced to the dynamical
        system defined by the Earth's mean equator and equinox, provide
        an epoch time.  To get J2000.0 coordinates, for example:

        >>> ra, dec, distance = ICRF([1, 2, 3]).radec(ts.J2000)
        >>> print(ra, dec, sep='\n')
        04h 13m 43.32s
        +53deg 17' 55.1"

        """
        position_au = self.position.au
        if epoch is not None:
            if isinstance(epoch, Time):
                pass
            elif isinstance(epoch, float):
                epoch = Time(None, tt=epoch)
            elif epoch == 'date':
                epoch = self.t
            else:
                raise ValueError('the epoch= must be a Time object,'
                                 ' a floating point Terrestrial Time (TT),'
                                 ' or the string "date" for epoch-of-date')
            position_au = einsum('ij...,j...->i...', epoch.M, position_au)
        r_au, dec, ra = to_polar(position_au)
        return (Angle(radians=ra, preference='hours'),
                Angle(radians=dec, signed=True),
                Distance(r_au))

    def separation_from(self, another_icrf):
        """Return the angle between this position and another.

        >>> print(ICRF([1,0,0]).separation_from(ICRF([1,1,0])))
        45deg 00' 00.0"

        You can also compute separations across an array of positions.

        >>> directions = ICRF([[1,0,-1,0], [0,1,0,-1], [0,0,0,0]])
        >>> directions.separation_from(ICRF([0,1,0])).degrees
        array([  90.,    0.,   90.,  180.])

        """
        p1 = self.position.au
        p2 = another_icrf.position.au
        u1 = p1 / length_of(p1)
        u2 = p2 / length_of(p2)
        if u2.ndim > 1:
            if u1.ndim == 1:
                u1 = u1[:,None]
        elif u1.ndim > 1:
            u2 = u2[:,None]
        c = dots(u1, u2)
        return Angle(radians=arccos(clip(c, -1.0, 1.0)))

    def ecliptic_position(self):
        """Compute J2000 ecliptic coordinates (x, y, z)"""
        vector = _ECLIPJ2000.dot(self.position.au)
        return Distance(vector)

    def ecliptic_latlon(self):
        """Compute J2000 ecliptic coordinates (lat, lon, distance)"""
        vector = _ECLIPJ2000.dot(self.position.au)
        d, lat, lon = to_polar(vector)
        return (Angle(radians=lat, signed=True),
                Angle(radians=lon),
                Distance(au=d))

    def galactic_position(self):
        """Compute galactic coordinates (x, y, z)"""
        vector = _GALACTIC.dot(self.position.au)
        return Distance(vector)

    def galactic_latlon(self):
        """Compute galactic coordinates (lat, lon, distance)"""
        vector = _GALACTIC.dot(self.position.au)
        d, lat, lon = to_polar(vector)
        return (Angle(radians=lat, signed=True),
                Angle(radians=lon),
                Distance(au=d))

    def to_skycoord(self, unit=None):
        """Convert this distance to an AstroPy ``SkyCoord`` object."""
        from astropy.coordinates import SkyCoord
        from astropy.units import au
        x, y, z = self.position.au
        return SkyCoord(representation='cartesian', x=x, y=y, z=z, unit=au)

    def _to_spice_frame(self, name):
        vector = self.position.au
        vector = inertial_frames[name].dot(vector)
        d, dec, ra = to_polar(vector)
        return (Angle(radians=ra, preference='hours', signed=True),
                Angle(radians=dec),
                Distance(au=d))

    def from_altaz(self, alt=None, az=None, alt_degrees=None, az_degrees=None):
        """Generate an Apparent position from an altitude and azimuth.

        The altitude and azimuth can each be provided as an `Angle`
        object, or else as a number of degrees provided as either a
        float or a tuple of degrees, arcminutes, and arcseconds::

            alt=Angle(...), az=Angle(...)
            alt_degrees=23.2289, az_degrees=142.1161
            alt_degrees=(23, 13, 44.1), az_degrees=(142, 6, 58.1)

        """
        # TODO: should this method live on another class?
        R = self.observer_data.altaz_rotation if self.observer_data else None
        if R is None:
            raise ValueError('only a position generated by a topos() call'
                             ' knows the orientation of the horizon'
                             ' and can understand altitude and azimuth')
        alt = _interpret_angle('alt', alt, alt_degrees)
        az = _interpret_angle('az', az, az_degrees)
        r = 0.1  # close enough to make gravitational refraction irrelevant
        p = from_polar(r, alt, az)
        p = einsum('ji...,j...->i...', R, p)
        return Apparent(p)

# For compatibility with my original name for the class.  Not an
# important enough change to warrant a deprecation error for users, so:
ICRS = ICRF

class Geometric(ICRF):
    """An (x,y,z) vector between two instantaneous position.

    A geometric position is the difference between the Solar System
    positions of two bodies at exactly the same instant.  It is *not*
    corrected for the fact that, in real physics, it will take time for
    light to travel from one position to the other.

    """
    def altaz(self, temperature_C=None, pressure_mbar='standard'):
        """Compute (alt, az, distance) relative to the observer's horizon

        The altitude returned is an `Angle` in degrees above the
        horizon, while the azimuth is the compass direction in degrees
        with north being 0 degrees and east being 90 degrees.

        """
        return _to_altaz(self.position.au, self.observer_data,
                         temperature_C, pressure_mbar)

class Barycentric(ICRF):
    """An (x, y, z) position measured from the Solar System barycenter.

    Each barycentric position is an ICRS position vector, meaning that
    the coordinate axes are defined by the high-precision ICRF that has
    replaced the old J2000.0 reference frame, and the coordinate origin
    is the BCRS gravitational center of the Solar System.

    Skyfield generates a `Barycentric` position whenever you ask a Solar
    System body for its location at a particular time:

    >>> t = ts.utc(2003, 8, 29)
    >>> mars.at(t)
    <Barycentric position and velocity at date t center=0 target=499>

    """
    def observe(self, body):
        """Compute the `Astrometric` position of a body from this location.

        To compute the body's astrometric position, it is first asked
        for its position at the time `t` of this position itself.  The
        distance to the body is then divided by the speed of light to
        find how long it takes its light to arrive.  Finally, the light
        travel time is subtracted from `t` and the body is asked for a
        series of increasingly exact positions to learn where it was
        when it emitted the light that is now reaching this position.

        >>> earth.at(t).observe(mars)
        <Astrometric position and velocity at date t>

        """
        p, v, light_time = body._observe_from_bcrs(self)
        t = self.t
        astrometric = Astrometric(p, v, t, observer_data=self.observer_data)
        astrometric.light_time = light_time
        return astrometric

# TODO: pre-create a Barycentric object representing the SSB, and make
# it possible for it to observe() a planet.

class Astrometric(ICRF):
    """An astrometric (x, y, z) position relative to a particular observer.

    The *astrometric position* of a body is its position relative to an
    observer, adjusted for light-time delay: the position of the body
    back when it emitted (or reflected) the light that is now reaching
    the observer's eyes or telescope.

    Astrometric positions are usually generated in Skyfield by calling
    the `Barycentric` method `observe()` to determine where a body will
    appear in the sky relative to a specific observer.

    """
    def apparent(self):
        """Compute an :class:`Apparent` position for this body.

        This applies two effects to the position that arise from
        relativity and shift slightly where the other body will appear
        in the sky: the deflection that the image will experience if its
        light passes close to large masses in the Solar System, and the
        aberration of light caused by the observer's own velocity.

        >>> earth.at(t).observe(mars).apparent()
        <Apparent position at date t>

        These transforms convert the position from the BCRS reference
        frame of the Solar System barycenter and to the reference frame
        of the observer.  In the specific case of an Earth observer, the
        output reference frame is the GCRS.

        """
        t = self.t
        position_au = self.position.au.copy()
        observer_data = self.observer_data
        gcrs_position = observer_data.gcrs_position

        if gcrs_position is None:
            include_earth_deflection = array((False,))
        else:
            limb_angle, nadir_angle = compute_limb_angle(
                position_au, gcrs_position)
            include_earth_deflection = nadir_angle >= 0.8

        add_deflection(position_au, observer_data.bcrs_position,
                       observer_data.ephemeris, t, include_earth_deflection)

        add_aberration(position_au, observer_data.bcrs_velocity,
                       self.light_time)

        return Apparent(position_au, t=t, observer_data=observer_data)

class Apparent(ICRF):
    """An apparent (x, y, z) position relative to a particular observer.

    The *apparent position* of a body is its position relative to an
    observer adjusted for light-time delay, deflection (light rays
    bending as they pass large masses like the Sun or Jupiter), and
    aberration (light slanting because of the observer's motion through
    space).

    Included in aberration is the relativistic transformation that takes
    the position out of the BCRS centered on the solar system barycenter
    and into the reference frame of the observer.  In the case of an
    Earth observer, the transform takes the coordinate into the GCRS.

    """
    def altaz(self, temperature_C=None, pressure_mbar='standard'):
        """Compute (alt, az, distance) relative to the observer's horizon

        The altitude returned is an `Angle` in degrees above the
        horizon, while the azimuth is the compass direction in degrees
        with north being 0 degrees and east being 90 degrees.

        """
        return _to_altaz(self.position.au, self.observer_data,
                         temperature_C, pressure_mbar)

class Geocentric(ICRF):
    """An (x,y,z) position measured from the geocenter."""

def _to_altaz(position_au, observer_data, temperature_C, pressure_mbar):
    """Compute (alt, az, distance) relative to the observer's horizon.

    """
    elevation_m = observer_data.elevation_m
    R = observer_data.altaz_rotation

    if (elevation_m is None) or (R is None):
        raise ValueError('to compute an altazimuth position, you must'
                         ' observe from a specific Earth location that'
                         ' you specify using a Topos instance')

    # TODO: wobble

    position_au = einsum('ij...,j...->i...', R, position_au)
    r_au, alt, az = to_polar(position_au)

    if temperature_C is None:
        alt = Angle(radians=alt)
    else:
        if temperature_C == 'standard':
            temperature_C = 10.0
        if pressure_mbar == 'standard':
            pressure_mbar = 1010.0 * exp(-elevation_m / 9.1e3)
        alt = refract(alt * RAD2DEG, temperature_C, pressure_mbar)
        alt = Angle(degrees=alt)

    return alt, Angle(radians=az), Distance(r_au)

def ITRF_to_GCRS(t, rITRF):  # todo: velocity

    # Todo: wobble

    spin = rot_z(t.gast * tau / 24.0)
    position = einsum('ij...,j...->i...', spin, array(rITRF))
    return einsum('ij...,j...->i...', t.MT, position)


# read snapshot and obtain multiple systems
import collections
from scipy import spatial as sp
from .base import *
from .bse import *

G_MSUN_PC_MYR=0.00449830997959438 # Msun, pc, myr
G_HENON=1 # Henon unit
HEADER_OFFSET=24 # header offset in bytes for snapshots with the BINARY format
HEADER_OFFSET_WITH_CM=72 # header offset with center-of-the-mass data in bytes for snapshots with the BINARY format

class PeTarDataHeader():
    """ Petar snapshot data header
    members:
        fid: int
           file id
        n: int
           number of particles
        time: float
           time of snapshot
        *pos_offset: list of float (length of 3)
           position offset of particle system
        *vel_offset: list of float (length of 3)
           velocity offset of particle system

        pos_offset and vel_offset only exist when keyword argument 'external_mode' is not none
    """

    def __init__(self, _filename=None, **kwargs):
        """ Initial data header

        Parameters:
        -----------
        _filename: string
            PeTar snapshot file name to read the header, if not provide, all members are initialized to zero (None)
        kwargs: dict
            Keyword arguments:
            snapshot_format: string (ascii)
                Data format of snapshot files: binary or ascii
            external_mode: string (none)
                PeTar external mode (set in configure): galpy, none
                If not none, this option indicates the pos_offset and vel_offset exists
        """
        self.fid = int(0)
        self.n = int(0)
        self.time = 0.0
        if ('external_mode' in kwargs.keys()):
            if (kwargs['external_mode']!='none'):
                self.pos_offset=[0.0,0.0,0.0]
                self.vel_offset=[0.0,0.0,0.0]

        if (_filename!=None): self.read(_filename,**kwargs)

    def read(self, _filename, **kwargs):
        """ Read snapshot file to obtain the header information

        Parameters:
        -----------
        _filename: string
            PeTar snapshot file name to read the header
        kwargs: dict
            Keyword arguments:
            snapshot_format: string (ascii)
                Data format of snapshot files: binary or ascii
            external_mode: string (none)
                PeTar external mode (set in configure): galpy, none
                If not none, this option indicates the pos_offset and vel_offset exists
        """
        snapshot_format='ascii'
        if ('snapshot_format' in kwargs.keys()): snapshot_format=kwargs['snapshot_format']
        offset_flag=False
        if ('external_mode' in kwargs.keys()):
            if (kwargs['external_mode']!='none'): offset_flag=True

        if (snapshot_format=='ascii'):
            fp = open(_filename, 'r')
            header=fp.readline()
            header_items=header.split()
            if (offset_flag):
                if (len(header_items)!=9):
                    raise ValueError('Snapshot header item number mismatch! Need 9 (file_id, N, time, xcm, ycm, zcm, vxcm, vycm, vzcm), got %d. Make sure the external_mode keyword set correctly.' % len(header_items))

                file_id, n_glb, t, x, y, z, vx, vy, vz = header_items
                fp.close()

                self.fid = int(file_id)
                self.n = int(n_glb)
                self.time = float(t)
                self.pos_offset = [float(x),float(y),float(z)]
                self.vel_offset = [float(vx),float(vy),float(vz)]
            else:
                if (len(header_items)!=3):
                    raise ValueError('Snapshot header item number mismatch! Need 3 (file_id, N, time), got %d. Make sure the external_mode keyword set correctly.' % len(header_items))

                file_id, n_glb, t = header_items
                fp.close()

                self.fid = int(file_id)
                self.n = int(n_glb)
                self.time = float(t)

        elif (snapshot_format=='binary'):
            if (offset_flag):
                fp = np.fromfile(_filename, dtype=np.dtype([('file_id',np.int64),('n_glb',np.int64),('time',np.float64),('x',np.float64),('y',np.float64),('z',np.float64),('vx',np.float64),('vy',np.float64),('vz',np.float64)]),count=1)
                self.file_id = fp['file_id'][0]
                self.n_glb = fp['n_glb'][0]
                self.time = fp['time'][0]
                self.pos_offset = [fp['x'][0], fp['y'][0], fp['z'][0]]
                self.vel_offset = [fp['vx'][0], fp['vy'][0], fp['vz'][0]]
            else:
                fp = np.fromfile(_filename, dtype=np.dtype([('file_id',np.int64),('n_glb',np.int64),('time',np.float64)]),count=1)
                self.file_id = fp['file_id'][0]
                self.n_glb = fp['n_glb'][0]
                self.time = fp['time'][0]
        else:
            raise ValueError('Snapshot format unknown, should be binary or ascii, given', snapshot_format)

class SimpleParticle(DictNpArrayMix):
    """ Simple particle class with only mass, postion, velocity
    keys: (class members)
        mass (1D): mass
        pos (2D,3): postion x, y, z
        vel (2D,3): velocity vx, vy, vz
    """
    def __init__(self, _dat=None, _offset=int(0), _append=False, **kwargs):
        """ DictNpArrayMix type initialzation, see help(DictNpArrayMix.__init__)
        """
        keys = [['mass', np.float64], ['pos', (np.float64, 3)], ['vel', (np.float64, 3)]]
        DictNpArrayMix.__init__(self, keys, _dat, _offset, _append, **kwargs)

    def calcR2(self):
        """ calculate distance square, r2, and add it as a class member
        """
        if (not 'r2' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['r2',1])
        self.r2 = vecDot(self.pos,self.pos)

    def calcEkin(self):
        """ calculate kinetic energy
        """
        if (not 'ekin' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['ekin',1])
        self.ekin = 0.5*vecDot(self.vel,self.vel)*self.mass

    def correctCenter(self, cm_pos, cm_vel):
        self.pos -= cm_pos
        self.vel -= cm_vel

class Particle(SimpleParticle):
    """ Particle class
        The particle data of PeTar. Depending on the compile configuration of PeTar,
        The data structures (columns) of the particle snapshots are different.
        Using the correct keyword arguments in the initialization to control the member definition (Keys)

    keys: (class members)
        The final keys are a combination of sub keys depending on keyword arguments (kwargs) of initial function

        Sub key list:

        std: [inherit SimpleParticle]
        bstat: binary_state: binary interruption state
        se: radius:        (1D): radius for merger checker
            dm:            (1D): mass loss
            time_record    (1D): last time of interruption check
            time_interrupt (1D): next interruption time
        bse: star  (SSEStarParameter): BSE based stellar evolution parameters
        ptcl: r_search (1D): searching radius
              id       (1D): identification
              mass_bk  (1D): artificial particle parameter 1
              status   (1D): artificial particle parameter 2
              r_in     (1D): changeover function inner boundary
              r_out    (1D): changeover function outer boundary
        hermite: dt    (1D): time step
                 time  (1D): current time
                 acc   (2D,3): acceleration x, y, z
                 jerk  (2D,3): acceleration derivative x, y, z
                 pot   (1D): potential
        soft: acc_soft (2D,3): long-range interaction acceleration (particle-tree) x, y, z
              pot      (1D): total potential
              pot_soft (1D): long-range interaction potential
              *pot_ext  (1D): external potential (only exist when keyword argument 'external_mode' is not 'none')
              n_nb:    (1D): number of neighbors (short-interaction)
        hermite: dt   (1D): hermite time step size
                 time (1D): current time of particle
                 acc  (2D,3): acceleration
                 jerk (2D,3): first derivates of acceleration
                 pot  (1D): potential

        ends: the end part of keys depends on kwargs['particle_type']:
             hermite:   ptcl + hermite
             hard:      ptcl
             soft (default): ptcl + soft

        final: the final combination of keys depends on kwargs['interrupt_mode']:
             base:      std + bstat + se + ends
             bse:       std + bstat + se + bse + ends
             none (default): std + bstat + ends

    """

    def __init__ (self, _dat=None, _offset=int(0), _append=False, **kwargs):
        """ DictNpArrayMix type initialzation, see help(DictNpArrayMix.__init__)

        Parameters
        ----------
        keyword arguments:
            particle_type: string (soft)
               basic particle type: hermite, hard, soft
            interrupt_mode: string (none)
               PeTar interrupt mode (set in configure): base, bse, mobse, none
               This option indicates whether columns of stellar evolution exist
            external_mode: string (none)
               PeTar external mode (set in configure): galpy, none
               This option indicates whether the column of externa potential exist
        """

        keys_bstat = [['binary_state',np.int64]]
        keys_se  = [['radius',np.float64],['dm',np.float64],['time_record',np.float64],['time_interrupt',np.float64]]
        keys_ptcl_add = [['r_search',np.float64], ['id',np.int64], ['mass_bk',np.int64], ['status',np.int64], ['r_in',np.float64], ['r_out',np.float64]]
        keys_hermite_add = [['dt',np.float64],['time',np.float64],['acc',(np.float64,3)],['jerk',(np.float64,3)],['pot',np.float64]]
        keys_soft_add = [['acc_soft',(np.float64,3)], ['pot',np.float64], ['pot_soft',np.float64], ['n_nb',np.int64]]
        if ('external_mode' in kwargs.keys()):
            if (kwargs['external_mode']!='none'):
                keys_soft_add = [['acc_soft',(np.float64,3)], ['pot',np.float64], ['pot_soft',np.float64], ['pot_ext',np.float64], ['n_nb',np.int64]]

        keys_end =  keys_ptcl_add + keys_soft_add
        if ('particle_type' in kwargs.keys()):
            if (kwargs['particle_type']=='hermite'):
                keys_end = keys_ptcl_add + keys_hermite_add
            elif (kwargs['particle_type']=='hard'):
                keys_end = keys_ptcl_add
        keys=keys_bstat+keys_end
        if ('interrupt_mode' in kwargs.keys()):
            if (kwargs['interrupt_mode']=='base'):
                keys = keys_bstat+keys_se+keys_end
            elif ('bse' in kwargs['interrupt_mode']):
                keys = keys_bstat+keys_se+[['star',SSEStarParameter]]+keys_end

        SimpleParticle.__init__(self, _dat, _offset, _append, **kwargs)
        DictNpArrayMix.__init__(self, keys, _dat, _offset+self.ncols, True, **kwargs)

    def calcEtot(self):
        """ Calculate total energy and add it as the member, etot
        """
        if (not 'etot' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['etot',np.float64])
        self.etot = self.ekin + self.mass*self.pot

def calculateParticleCMDict(pcm, _p1, _p2):
    """ Calculate the center-of-the-mass of two particle sets

    Parameters
    ----------
    _p1: inherited SimpleParticle
        particle set 1
    _p2: inherited SimpleParticle
        particle set 2, should have the same size as _p1
    pcm: dict
        particle center-of-the-mass, should include keys: 'mass','pos','vel'.
    """
    if (issubclass(type(_p1), SimpleParticle)) & (issubclass(type(_p2),SimpleParticle)):
        pcm['mass'] = _p1.mass + _p2.mass
        pcm['pos']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1.mass, _p1.pos, _p2.mass, _p2.pos)))
        pcm['vel']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1.mass, _p1.vel, _p2.mass, _p2.vel)))
    elif (isinstance(_p1, collections.OrderedDict)) & (isinstance(_p2,collections.OrderedDict)) | (isinstance(_p1, dict)) & (isinstance(_p2, dict)):
        pcm['mass'] = _p1['mass'] + _p2['mass']
        pcm['pos']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1['mass'], _p1['pos'], _p2['mass'], _p2['pos'])))
        pcm['vel']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1['mass'], _p1['vel'], _p2['mass'], _p2['vel'])))
    else:
        raise ValueError('Initial fail, date type should be Particle or collections.OrderDict, given',type(_p1))

class Binary(SimpleParticle):
    """ Binary class
        The binary (tree) data. Depending on the definition of two members
        (keyword argument member_particle_type(|_one|_two),
        The binary can refer to any type of multiple system.

    Keys:
        The final keys depends on kwargs of initial function

        kwargs['simple_mode'] (bool)
            True: (default)
                mass (1D): total mass of two components
                pos  (2D,3): c.m. position x, y, z
                vel  (2D,3): c.m. velocity vx, vy, vz
                rrel (1D): relative distance
                semi (1D): semi-major axis
                ecc  (1D): eccentricity
                p1   (member_particle_type_one) component one
                p2   (member_particle_type_two) component two
            False:
                mass (1D): total mass of two components
                pos  (2D,3): c.m. position x, y, z
                vel  (2D,3): c.m. velocity vx, vy, vz
                m1   (1D): component 1 mass
                m2   (1D): component 2 mass
                rrel (1D): relative distance
                semi (1D): semi-major axis
                am   (2D,3): specific angular momemtum x, y, z
                L    (2D,3): angular momemtum x, y, z
                eccvec  (2D,3): eccentric vector
                incline (1D): inclination
                rot_horizon (1D): frame rotational angle in x-y plane (longitude of ascending node)
                ecc  (1D): eccentricity
                rot_self (1D): frame rotational angle in orbital plane (argument of periapsis)
                ecca (1D): eccentric anomaly
                period (1D): period
                t_peri (1D): time to peri-center
                p1 (member_particle_type_one) component one
                p2 (member_particle_type_two) component two

        The member_particle_type(|_one|_two) is given by keyword arguments:
           'member_particle_type' (for both members),'member_particle_type_one','member_particle_type_two'.
        In default, it is petar.SimpleParticle.
        If a type (e.g., petar.Particle) is given, the member is a single star.
        If a list with two members (e.g., [petar.Particle, petar.Particle]) is given,
        the member is a binary with two single stars.
        A hierarchical list can be provided, e.g., [petar.Particle, [petar.Particle, petar.Particle]]
        to indicate a triple system.

    """
    def __init__ (self, _p1=None, _p2=None, _offset=int(0), _append=False, **kwargs):
        """
        Parameters
        ----------
        _p1: inherited SimpleParticle | 2D numpy.ndarray | Binary | None
            If the type is inherited SimpleParticle, it is the first component of binary (_p2 should be the same type).
            If the type is Binary, the class instance is initialized by copy the data of _p1.
            If it is None, initialize class with empty data
        _p2: inherited SimpleParticle | None
            If the type is inherited SimpleParticle, it is the second component of binary
            If it is None, _p1 should be either 2D numpy.ndarray or Bina
        _offset: int (0)
            Reading column offset of _dat if it is 2D np.ndarray
        _append: bool (False)
            If true, append keys and ncols to the current class instead of create new class members

        keyword arguments:
            simple_mode: bool (True)
                If True, only calculate semi and ecc, save computing time significantly
            G: float (1.0)
                Gravitational constant
            member_particle_type: type or list (SimpleParticle)
                Type of component particle (both)
            member_particle_type_one: type or list (SimpleParticle)
                Type of 1st component
            member_particle_type_two: type or list (SimpleParticle)
                Type of 2nd component
        """
        G=1
        simple_mode=True
        member_particle_type=SimpleParticle
        member_particle_type_one=member_particle_type
        member_particle_type_two=member_particle_type

        if 'G' in kwargs.keys(): G=kwargs['G']
        if 'simple_mode' in kwargs.keys(): simple_mode=kwargs['simple_mode']
        if 'member_particle_type' in kwargs.keys():
            member_particle_type=kwargs['member_particle_type']
            member_particle_type_one=member_particle_type
            member_particle_type_two=member_particle_type
        if 'member_particle_type_one' in kwargs.keys(): member_particle_type_one=kwargs['member_particle_type_one']
        if 'member_particle_type_two' in kwargs.keys(): member_particle_type_two=kwargs['member_particle_type_two']

        if (issubclass(type(_p1), SimpleParticle)) & (issubclass(type(_p2),SimpleParticle)):
            if (simple_mode):
                self.keys = [['mass',np.float64],['pos',(np.float64,3)],['vel',(np.float64,3)],['rrel',np.float64],['semi',np.float64],['ecc',np.float64],['p1',(type(_p1),_p1.initargs)], ['p2', (type(_p2),_p2.initargs)]]
                self.particleToSemiEcc(_p1, _p2, G)
                self.ncols= int(10)
            else:
                self.keys = [['mass',np.float64],['pos',(np.float64,3)],['vel',(np.float64,3)],['m1',np.float64],['m2',np.float64],['rrel',np.float64],['semi',np.float64],['am',(np.float64,3)],['L',(np.float64,3)],['eccvec',(np.float64,3)],['incline',np.float64],['rot_horizon',np.float64],['ecc',np.float64],['rot_self',np.float64],['ecca',np.float64],['period',np.float64],['t_peri',np.float64],['p1',(type(_p1),_p1.initargs)], ['p2', (type(_p2),_p2.initargs)]]
                self.particleToBinary(_p1, _p2, G)
                self.ncols= int(27)
            self.p1 = _p1
            self.p2 = _p2
            self.size = _p1.size
            self.ncols += self.p1.ncols + self.p2.ncols
            self.initargs = kwargs.copy()
            binary_tree = self.createMemberParticleTypeTree()
            self.initargs['member_particle_type_one']=binary_tree[0]
            self.initargs['member_particle_type_two']=binary_tree[1]
        elif (_p2==None):
            type_one = member_particle_type_one
            if (type(member_particle_type_one) == list):
                type_one = (Binary, {'member_particle_type_one':member_particle_type_one[0],'member_particle_type_two':member_particle_type_one[1]})
            type_two = member_particle_type_two
            if (type(member_particle_type_two) == list):
                type_two = (Binary, {'member_particle_type_one':member_particle_type_two[0],'member_particle_type_two':member_particle_type_two[1]})
            if (simple_mode):
                keys = [['rrel',np.float64],['semi',np.float64],['ecc',np.float64],['p1',type_one], ['p2', type_two]]
                SimpleParticle.__init__(self, _p1, _offset, _append, **kwargs)
                DictNpArrayMix.__init__(self, keys, _p1, _offset+self.ncols, True, **kwargs)
            else:
                keys=[['m1',np.float64],['m2',np.float64],['rrel',np.float64],['semi',np.float64],['am',(np.float64,3)],['L',(np.float64,3)],['eccvec',(np.float64,3)],['incline',np.float64],['rot_horizon',np.float64],['ecc',np.float64],['rot_self',np.float64],['ecca',np.float64],['period',np.float64],['t_peri',np.float64],['p1', type_one],['p2', type_two]]
                SimpleParticle.__init__(self, _p1, _offset, _append, **kwargs)
                DictNpArrayMix.__init__(self, keys, _p1, _offset+self.ncols, True, **kwargs)
            self.initargs = kwargs.copy()
        else:
            raise ValueError('Initial fail, date type should be Particle (2), Binary (1) or no argument (0)')

    def calcEkin(self):
        """ Calculate c.m. kinetic energy, ekin, and add it as a member
        """
        ekin = 0.5*vecDot(self.vel,self.vel)*self.mass
        self.addNewMember('ekin',ekin)

    def calcEtot(self):
        """ Calculate c.m. total energy (binary energy is excluded) , etot, and add it as a member
        """
        etot = self.ekin + self.mass*self.pot
        self.addNewMember('etot',etot)

    def calcR2(self, member_also=False):
        """ Calculate c.m. distance square, r2, and add it as a member
        """
        r2 = vecDot(self.pos,self.pos)
        self.addNewMember('r2',r2)
        if (member_also):
            ncols = self.p1.ncols + self.p2.ncols
            self.p1.calcR2()
            self.p2.calcR2()
            ncols = self.p1.ncols + self.p2.ncols - ncols
            self.ncols += ncols

    def calcEbin(self):
        """ Calculate binding energy, ebin, and add it as a member
            Notice G should be given the correct value in initialization (keyword argument 'G')
        """
        ebin = self.initargs['G']*self.p1.mass*self.p2.mass/(2*self.semi)
        self.addNewMember('ebin',ebin)

    def calcPot(self):
        """ Calculate potential of c.m., pot, and add it as a member
            Notice G should be given the correct value in initialization (keyword argument 'G')
        """
        G = self.initargs['G']
        pos_b1 = self.p1.pos
        pos_b2 = self.p2.pos
        m_b1 = self.p1.mass
        m_b2 = self.p2.mass
        dr = pos_b1-pos_b2
        dr2 = vecDot(dr,dr)
        invr = 1/np.sqrt(dr2)
        pot_b1 = self.p1.pot + G*m_b2*invr
        pot_b2 = self.p2.pot + G*m_b1*invr
        pot = (m_b2*pot_b1 + m_b1*pot_b2)/self.mass
        self.addNewMember('pot',pot)

    def generateBinaryID(self):
        """ Use CantorPairing to map two components id to one binary id
            Add new member bid
        """
        bid = cantorPairing(self.p1.id, self.p2.id)
        self.addNewMember('bid',bid)

    def correctCenter(self, cm_pos, cm_vel):
        """ Corrent c.m and component position and velocity by subtracting cm_pos and cm_vel
        """
        self.pos -= cm_pos
        self.vel -= cm_vel
        self.p1.correctCenter(cm_pos, cm_vel)
        self.p2.correctCenter(cm_pos, cm_vel)

    def particleToSemiEcc(self, _p1, _p2, _G):
        """ Calculate relative distance, semi-major axis and eccentricity from particle pairs

        Parameters
        ----------
        _p1, _p2: inherited SimpleParticle
            Particle pair data set
        _G: float
            Gravitational constant

        """
        calculateParticleCMDict(self.__dict__, _p1, _p2)

        dr = (_p1.pos - _p2.pos)
        dv = (_p1.vel - _p2.vel)

        dr2  = (dr*dr).sum(axis=1)
        dv2  = (dv*dv).sum(axis=1)
        rvdot= (dr*dv).sum(axis=1)

        dr   = np.sqrt(dr2)
        m    = (_p1.mass+_p2.mass)
        semi = 1.0/(2.0/dr - dv2/(_G*m))

        dr_semi = 1.0 - dr/semi
        ecc = np.sqrt(dr_semi*dr_semi + rvdot*rvdot/(_G*m*semi))

        self.rrel = dr
        self.semi = semi
        self.ecc  = ecc

    def particleToBinary(self, _p1, _p2, _G):
        """ Calculate binary orbit from particle pairs

        Parameters
        ----------
        _p1, _p2: inherited SimpleParticle
            Particle pair data set
        _G: float
            Gravitational constant

        """
        binary=self.__dict__

        def regular_sign(_a,_a_err):
            _a[(_a<0) & (_a>-_a_err)] *= -1

        f_err = 1e-2
        calculateParticleCMDict(binary, _p1, _p2)

        binary['m1'] = _p1.mass
        binary['m2'] = _p2.mass
        m_tot = binary['mass']
        Gm_tot = _G*m_tot

        dx = _p1.pos-_p2.pos
        dv = _p1.vel-_p2.vel
        dr2  = vecDot(dx,dx)
        dv2  = vecDot(dv,dv)
        rvdot= vecDot(dx,dv)
        dr   = np.sqrt(dr2)
        binary['rrel'] = np.sqrt(dr2)

        inv_dr = 1.0 / binary['rrel']
        binary['semi'] = 1.0 / (2.0*inv_dr - dv2 / Gm_tot)
        binary['am'] = np.cross(dx,dv)
        dp = _p1.vel*_p1.mass[:,None] - _p2.vel*_p2.mass[:,None]
        binary['L'] = np.cross(dx,dp)
        binary['eccvec'] = np.cross(dv,binary['am'])/Gm_tot[:,None]-dx/dr[:,None]

        binary['incline'] = np.arctan2(np.sqrt(binary['am'][:,0]*binary['am'][:,0]+binary['am'][:,1]*binary['am'][:,1]),binary['am'][:,2])
        binary['rot_horizon'] = np.arctan2(binary['am'][:,0],-binary['am'][:,1])
        regular_sign(binary['am'][:,0],f_err)
        regular_sign(binary['am'][:,1],f_err)
        #binary['rot_horizon'][binary['rot_horizon']<0] += np.pi
        binary['rot_horizon'][binary['am'][:,1]==0.0]=0.0

        cosOMG = np.cos(binary['rot_horizon'])
        sinOMG = np.sin(binary['rot_horizon'])
        cosinc = np.cos(binary['incline'])
        sininc = np.sin(binary['incline'])

        pos_bar_x =   dx[:,0]*cosOMG + dx[:,1]*sinOMG
        pos_bar_y = (-dx[:,0]*sinOMG + dx[:,1]*cosOMG)*cosinc + dx[:,2]*sininc
        pos_bar_z = 0.0
        vel_bar_x =   dv[:,0]*cosOMG + dv[:,1]*sinOMG
        vel_bar_y = (-dv[:,0]*sinOMG + dv[:,1]*cosOMG)*cosinc + dv[:,2]*sininc
        vel_bar_z = 0.0

        h = np.sqrt(np.sum(binary['am']*binary['am'],axis=1))
        ecccosomg =  h/Gm_tot*vel_bar_y - pos_bar_x*inv_dr
        eccsinomg = -h/Gm_tot*vel_bar_x - pos_bar_y*inv_dr
        binary['ecc'] = np.sqrt( ecccosomg*ecccosomg + eccsinomg*eccsinomg )
        regular_sign(ecccosomg,f_err)
        regular_sign(eccsinomg,f_err)
        binary['rot_self'] = np.arctan2(eccsinomg,ecccosomg)
        #binary['rot_self'][binary['rot_self']<-np.pi+1e-5] += 2*np.pi
        #binary['rot_self'][binary['rot_self']>=np.pi-1e-5] -= 2*np.pi

        regular_sign(pos_bar_y,f_err)
        regular_sign(pos_bar_x,f_err)
        phi = np.arctan2(pos_bar_y, pos_bar_x)
        #phi[phi<-np.pi+1e-5] += 2*np.pi
        #phi[phi>=np.pi-1e-5] -= 2*np.pi

        f = phi - binary['rot_self']
        binary['ecca'] = np.arctan(np.sin(f)*np.sqrt(np.abs(binary['ecc']*binary['ecc'] - 1.0))/(binary['ecc']+np.cos(f)))
        n = np.sqrt(Gm_tot/np.abs(binary['semi']*binary['semi']*binary['semi']))
        binary['period'] = 8.0*np.arctan(1.0)/n
        l = binary['ecca'] - binary['ecc']*np.sin(binary['ecca'])
        binary['t_peri'] = l / n

    def createMemberParticleTypeTree(self):
        """ scan the members to create the member particle type tree list
            For example, if the binary structure is a triple: p1: single, p2: binary.
            Then the returned tree list is [particle_typename, [particle_typename, particle_typename]]
        """
        binary_tree=[None,None]
        if (type(self.p1) == Binary):
            binary_tree[0] = self.p1.createMemberParticleTypeTree()
        else:
            binary_tree[0] = type(self.p1)
        if (type(self.p2) == Binary):
            binary_tree[1] = self.p2.createMemberParticleTypeTree()
        else:
            binary_tree[1] = type(self.p2)
        return binary_tree

def findPair(_dat, _G, _rmax, use_kdtree=False, simple_binary=True):
    """  Find binaries in a particle data set
    The scipy.spatial.cKDTree is used to find pairs

    Parameters
    ----------
    _dat: inhermited SimpleParticle
        Particle data set
    _G: float
        Gravitational constant
    _rmax: float
        Maximum binary separation
    use_kdtree: bool (False)
        If True, use KDtree to find all binaries (slow); otherwise use information from PeTar, only hard binaries are detected (fast)
    simple_binary: bool (True)
        If True, only calculate semi and ecc (fast); otherwise calculating all binary parameters (slow)

    Return
    ----------
    kdt: KDtree structure if use_kdtree=True
    single: type of _dat
        single particle data set
    binary: Binary(simple_mode=simple_binary, member_particle_type=type(single), G=_G)
        binary data set
    """
    if (not issubclass(type(_dat), SimpleParticle)):
        raise ValueError("Data type wrong",type(_dat)," should be subclass of ", SimpleParticle)

    if (use_kdtree):
        # create KDTree
        #print('create KDTree')
        kdt=sp.cKDTree(_dat.pos)

        # find all close pairs
        #pairs=kdt.query_pairs(_rmax*AU2PC)

        # only check nearest index
        #pair_index=np.unique(np.transpose(np.array([np.array([x[0],x[1]]) for x in pairs])),axis=0)

        # find pair index and distance
        #print('Get index')
        r,index=kdt.query(_dat.pos,k=2)
        pair_index=np.transpose(np.unique(np.sort(index,axis=1),axis=0))
        #pair_index = np.transpose(index)

        #index = kdt.query_pairs(_rmax,output_type='ndarray')
        #pair_index = np.transpose(index)

        # two members
        p1 = _dat[pair_index[0]]
        p2 = _dat[pair_index[1]]

        # check orbits
        #print('Create binary')
        binary = Binary(p1, p2, G=_G, simple_mode=simple_binary)
        apo =binary.semi*(binary.ecc+1.0)

        bsel= ((binary.semi>0) & (apo<_rmax))
        binary = binary[bsel]

        single_mask = np.ones(_dat.size).astype(bool)
        single_mask[pair_index[0][bsel]]=False
        single_mask[pair_index[1][bsel]]=False
        single = _dat[single_mask]
        return kdt, single, binary
    else:
        idx = _dat.status.argsort()
        dat_sort = _dat[idx]
        status, index, inverse, counts = np.unique(dat_sort.status, return_index=True, return_inverse=True, return_counts=True)
        binary_i1 = index[counts==2]
        binary_i2 = binary_i1+1
        binary = Binary(dat_sort[binary_i1], dat_sort[binary_i2], _G)
        single = dat_sort[index[-1]:]

        return single, binary

def findMultiple(_single, _binary, _G, _rmax, simple_binary=True):
    """  Find triples and quadruples from single and binary data
    The scipy.spatial.cKDTree is used to find pairs

    Parameters
    ----------
    _single: inhermited SimpleParticle
        Single particle data set
    _binary: Binary
        Binary data set
    _G: float
        Gravitational constant
    _rmax: float
        Maximum binary separation
    simple_binary: bool (True)
        If True, only calculate semi and ecc (fast); otherwise calculating all binary parameters (slow)

    Return
    ----------
    kdt: KDtree structure if use_kdtree=True
    single: type of _dat
        single particle data set
    binary: Binary(simple_mode=simple_binary, member_particle_type=type(single), G=_G)
        binary data set
    triple: Binary(p1: type(single), p2: type(binary), G=_G)
        triple data set
    quadruple: Binary(p1: type(binary), p2: type(binary), G=_G)
        quadruple (binary-binary) data set
    """
    if (not issubclass(type(_single), SimpleParticle)):
        raise ValueError("Data type wrong",type(_single)," should be subclass of ", SimpleParticle)

    single_sin = SimpleParticle(_single)
    binary_sin = SimpleParticle(_binary)
    all_sin = join(single_sin, binary_sin)

    # create KDTree
    kdt=sp.cKDTree(all_sin.pos)

    # find pair index and distance
    r,index=kdt.query(all_sin.pos,k=2)
    pair_index=np.transpose(np.unique(np.sort(index,axis=1),axis=0))

    bout_i1 = pair_index[0]
    bout_i2 = pair_index[1]

    Ns = _single.size
    Nb = _binary.size
    quad_pre_sel= (bout_i1>=Ns) & (bout_i2>=Ns)
    tri_pre_sel = (bout_i1<Ns) & (bout_i2>=Ns)
    bin_pre_sel = (bout_i1<Ns) & (bout_i2<Ns)

    n_quad_pre = quad_pre_sel.sum()
    n_tri_pre = tri_pre_sel.sum()
    n_bin_pre = bin_pre_sel.sum()
    if (bout_i1.size != n_quad_pre + n_tri_pre + n_bin_pre):
        raise ValueError('Error: multiple index selection size miss match: dat:',bout_i1.size,'quad:',n_quad_pre,'tri:',n_tri_pre,'bin:',n_bin_pre)

    s_del_index=np.array([]).astype(int)
    b_del_index=np.array([]).astype(int)

    quadruple = Binary(member_particle_type = [type(_single), type(_single)], **{**_single.initargs, 'G':_G, 'simple_mode':simple_binary})
    if (quad_pre_sel.sum()):
        q1_index = bout_i1[quad_pre_sel]-Ns
        q2_index = bout_i2[quad_pre_sel]-Ns
        quad_pre = Binary(_binary[q1_index], _binary[q2_index], **{**_single.initargs, 'G':_G, 'simple_mode':simple_binary})
        apo = quad_pre.semi*(quad_pre.ecc+1.0)
        quad_sel = (quad_pre.semi>0) & (apo<_rmax)
        quadruple = quad_pre[quad_sel]
        b_del_index=np.append(q1_index[quad_sel],q2_index[quad_sel])

    triple = Binary(member_particle_type_one = type(_single),
                    member_particle_type_two = [type(_single), type(_single)],
                    **{**_single.initargs, 'G':_G, 'simple_mode':simple_binary})
    if (tri_pre_sel.sum()):
        s_index = bout_i1[tri_pre_sel]
        b_index = bout_i2[tri_pre_sel]-Ns
        tri_pre = Binary(_single[s_index], _binary[b_index], **{**_single.initargs, 'G':_G, 'simple_mode':simple_binary})
        apo = tri_pre.semi*(tri_pre.ecc+1.0)
        tri_sel = (tri_pre.semi>0) & (apo<_rmax)
        triple = tri_pre[tri_sel]
        b_del_index=np.append(b_del_index,b_index[tri_sel])
        s_del_index=s_index[tri_sel]

    bmask=np.ones(Nb).astype(bool)
    if (b_del_index.size>0): bmask[b_del_index]=False;
    binary = _binary[bmask]

    if (bin_pre_sel.sum()):
        s1_index = bout_i1[bin_pre_sel]
        s2_index = bout_i2[bin_pre_sel]
        bin_pre = Binary(_single[s1_index], _single[s2_index], **{**_single.initargs, 'G':_G, 'simple_mode':simple_binary})
        apo = bin_pre.semi*(bin_pre.ecc+1.0)
        bin_sel = (bin_pre.semi>0) & (apo<_rmax)
        binary.append(bin_pre[bin_sel])
        s_del_index = np.concatenate((s_del_index, s1_index[bin_sel], s2_index[bin_sel]))

    smask=np.ones(Ns).astype(bool)
    smask[s_del_index]=False
    single = _single[smask]

    return single, binary, triple, quadruple


import os
import sys
import subprocess
import datetime
import numpy as np
from astropy.io import ascii
from astropy.table import Table, Column, vstack
from astropy.wcs import WCS
from astropy.wcs.utils import proj_plane_pixel_scales
from astropy.coordinates import SkyCoord, ICRS
from astropy.stats import gaussian_fwhm_to_sigma, sigma_clipped_stats
from scipy.stats import norm, f
from scipy.odr import *
from scipy.optimize import minimize
from scipy.ndimage.filters import median_filter, gaussian_filter1d
from photutils import detect_sources, Background

# For debugging
import matplotlib.pyplot as plt
import pdb

# Add the AstroImage class
sys.path.append("C:\\Users\\Jordan\\Libraries\\python\\AstroImage")
from AstroImage import AstroImage
import image_tools

# This script will read in the background level estimated for each on-target
# image in the previous step. The background level in dimmest parts of the
# on-target image will be directly computed, and the residual between the direct
# estimate and the interpolation will be stored. The distribution of these
# residual will be used to estimate which interpolated background levels can be
# trusted.

#==============================================================================
# *********************** CUSTOM USER CODE ************************************
# this is where the user specifies where the raw data is stored
# and some of the subdirectory structure to find the actual .FITS images
#==============================================================================

# This is the location of all PPOL reduction directory
PPOL_dir = 'C:\\Users\\Jordan\\FITS_data\\Mimir_data\\PPOL_reduced'

# Build the path to the S3_Asotremtry files
S3dir = os.path.join(PPOL_dir, 'S3_Astrometry')

# This is the location where all pyPol data will be saved
pyPol_data = 'C:\\Users\\Jordan\\FITS_data\\Mimir_data\\pyPol_data'

# This is the directory where the 2MASS tiles of the targets have been saved
# Go to "http://hachi.ipac.caltech.edu/" to download 2MASS tiles
TMASSdir  = "C:\\Users\\Jordan\\Libraries\\python\\Mimir_pyPol\\2MASSimages"

# Setup new directory for background subtracted data
bkgSubDir = os.path.join(pyPol_data, 'bkgSubtracted')
if (not os.path.isdir(bkgSubDir)):
    os.mkdir(bkgSubDir, 0o755)

# Read in Kokopelli mask generated in previous step
kokopelliMask = (AstroImage('kokopelliMask.fits').arr != 0)

# Read in the indexFile data and select the filenames
indexFile = os.path.join(pyPol_data, 'reducedFileIndex.csv')
fileIndex = Table.read(indexFile, format='csv')

# Grab the file basenames for later use
fileIndexFileNames = np.array([os.path.basename(file1)
    for file1 in fileIndex['Filename'].data])

# Modify the fileIndex to include rejections by residual value
if 'Background Cut' not in fileIndex.keys():
    fileIndex.add_column(Column(name='Background Cut',
                                data = np.repeat(0, len(fileIndex))))

# Determine which parts of the fileIndex pertain to science images
useFiles = np.where(np.logical_and(fileIndex['Use'].data == 1,
                                   fileIndex['Background'].data >= 0))
skipFiles = np.where(np.logical_or(fileIndex['Use'].data == 0,
                                   fileIndex['Background'].data < 0))

# Cull the file index to only include files selected for use
fileIndex1 = fileIndex[useFiles]
fileIndex2 = fileIndex[skipFiles]

# Group files by target and waveband
groupFileIndex = fileIndex1.group_by(['PPOL Name'])

allFileList     = []
allResidualList = []
# Loop through all the usable images and comute their residuals
for group in groupFileIndex.groups:
    # Grab the current target information
    thisTarget   = str(np.unique(group['Target'].data)[0])
    thisWaveband = str(np.unique(group['Waveband'].data)[0])
    thisPPOLname = str(np.unique(group['PPOL Name'].data)[0])

    # if thisPPOLname != 'NGC2023_H3': continue

    print('\nProcessing images for')
    print('\tPPOL Group : {0}'.format(thisPPOLname))
    print('')

    # Read in the 2MASS image
    TMASSfile = os.path.join(TMASSdir, '_'.join([thisTarget, thisWaveband]) + '.fits')
    TMASSimg  = AstroImage(TMASSfile)
    TMASSwcs  = WCS(TMASSimg.header)

    # Estimate the "nebula free" level
    mean, median, stddev = sigma_clipped_stats(TMASSimg.arr.flatten())
    bkgThresh = median - 0.5*stddev

    # Find the "nebula free" pixels
    bkgRegion = TMASSimg.arr < bkgThresh
    neighborCount = np.zeros_like(bkgRegion, dtype=int)
    for dx in range(-1,2):
        for dy in range(-1,2):
            neighborCount += np.roll(np.roll(bkgRegion, dy, axis = 0), dx, axis = 1)

    # Find pixels with at least 3 neighbors (other than self)
    bkgRegion = neighborCount > 4

    groupFileList     = []
    groupResidualList = []
    for file1, interpBkg in zip(group['Filename'].data, group['Background'].data):
        # Read in this image.
        img = AstroImage(file1)

        # See which pixels in this image map to background pixels
        ny, nx = img.arr.shape
        yy, xx = np.mgrid[0:ny, 0:nx]
        wcs = WCS(img.header)
        RAs, Decs = wcs.wcs_pix2world(xx, yy, 0)
        Tx, Ty    = TMASSwcs.wcs_world2pix(RAs, Decs, 0)
        Tx, Ty    = (Tx.round()).astype(int), (Ty.round()).astype(int)

        # Grab the value of the TMASS background mask for each pixel
        MimirBkgRegion = bkgRegion[Ty, Tx]

        # Get the indices of the background pixel
        bkgInds  = np.where(MimirBkgRegion)
        bkgVals  = img.arr[bkgInds]

        # Compute the direct estimate of background level
        mean, median, stddev = sigma_clipped_stats(bkgVals)

        # Compute the residual level and store it in the list
        thisResidual = mean - interpBkg
        groupFileList.append(os.path.basename(file1))
        groupResidualList.append(thisResidual)

    # Place this residual list in the final total residual list
    allFileList.extend(groupFileList)
    allResidualList.extend(groupResidualList)

    # Convert the lists to arrays
    groupFileList     = np.array(groupFileList)
    groupResidualList = np.array(groupResidualList)

    # Check for outliers and mark residuals 5-sigma outside this group's median
    mean, median, stddev = sigma_clipped_stats(groupResidualList)
    residMin, residMax = mean - 5*stddev, mean + 5*stddev
    badInds = np.where(np.logical_or(groupResidualList < residMin,
                                     groupResidualList > residMax))

    # If some of these residuals are more than 5-sigma from the group mean, then
    # mark them as bad background levels in the file index.
    if len(badInds[0]) > 0:
        # Select the file names of the bad backgrounds
        badFiles = groupFileList[badInds]
        # Grab the indices of these files in the fileIndex and mark them as bad
        fileIndexInds = np.array([np.where(fileIndexFileNames == file1)[0][0]
            for file1 in badFiles])
        fileIndex['Background Cut'][fileIndexInds] = 1

# Convert the lists to arrays
allFileList     = np.array(allFileList)
allResidualList = np.array(allResidualList)

# Now that we have the residuals for each group, plot them up as histograms
# # Start by parsing out the residuals for each group
# Now create a plot with all groups clumpped together
fig2 = plt.figure()
ax2  = fig2.add_subplot(1,1,1)
ax2.hist(allResidualList, 10, normed=1, histtype='stepfilled', stacked=True)
plt.xlabel('Residual Counts')
plt.ylabel('Fraction of Fields')

# Prepare some statistical comments
xmin, xmax = ax2.get_xlim()
ymin, ymax = ax2.get_ylim()
mean, median, stddev = sigma_clipped_stats(allResidualList)

# Mark the mean
ax2.axvline(mean, color='k', linewidth=2.0)
ax2.text(mean+0.02*(xmax-xmin), 0.95*ymax, 'mean', rotation='vertical')

# Mark the median
ax2.axvline(median, color='k', linewidth=2.0)
ax2.text(median-0.04*(xmax-xmin), 0.95*ymax, 'median', rotation='vertical')

# Mark the 3-sigma upper and lower limits
ax2.axvline(median - 5*stddev, color='k', linewidth=2.0)
ax2.axvline(median + 5*stddev, color='k', linewidth=2.0)

# Prepare the limits of the acceptable residual range
residMin, residMax = mean - 5*stddev, mean + 5*stddev

# Find any background levels that are outside the 5-sigma limits
badInds = np.where(np.logical_or(allResidualList < residMin,
                                 allResidualList > residMax))

# If some of these residuals are more than 5-sigma from the group mean, then
# mark them as bad background levels in the file index.
if len(badInds[0]) > 0:
    # Select the file names of the bad backgrounds
    badFiles = allFileList[badInds]
    # Grab the indices of these files in the fileIndex and mark them as bad
    fileIndexInds = np.array([np.where(fileIndexFileNames == file1)[0][0]
        for file1 in badFiles])
    fileIndex['Background Cut'][fileIndexInds] = 1

# Then save to disk
print('*************************************')
print('Writing all background levels to disk')
print('*************************************')
pdb.set_trace()
fileIndex.write(indexFile, format='csv')

print('Done!')


#!/usr/bin/env python3
#    Copyright (c) 2021 Michele Mancarella <michele.mancarella@unige.ch>
#
#    All rights reserved. Use of this source code is governed by a modified BSD
#    license that can be found in the LICENSE file.

from .ABSdata import Data, LVCData

import numpy as np
import astropy.units as u
import h5py
import os
import sys
#from pesummary.io import read
#import glob

PACKAGE_PARENT = '..'
SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))

from astropy.cosmology import Planck15
from cosmology.cosmo import Cosmo

import Globals

class O1O2Data(LVCData):

    def __init__(self, fname, which_metadata='GWOSC',  **kwargs):#nObsUse=None, nSamplesUse=None, dist_unit=u.Gpc, events_use=None, which_spins='skip' ):

        self.post_file_extension='.hdf5'
        import pandas as pd
        if which_metadata=='GWOSC':
            print('Using SNRS and far from the public version of the GWTC-3 catalog from the GWOSC')
            self.metadata = pd.read_csv(os.path.join(fname, 'GWTC-1-confident.csv'))
        else:
            print('Using best SNRS and far from all pipelines as reported in the GWTC-3 catalog paper')
            self.metadata = pd.read_csv(os.path.join(Globals.dataPath, 'all_metadata_pipelines_best.csv'))

        LVCData.__init__(self, fname, **kwargs) #nObsUse=nObsUse, nSamplesUse=nSamplesUse, dist_unit=dist_unit, events_use=events_use, which_spins=which_spins)

    def _set_Tobs(self):
        # The first observing run (O1) ran from September 12th, 2015 to January 19th, 2016 --> 129 days
        # From https://journals.aps.org/prx/pdf/10.1103/PhysRevX.6.041015:
        # after data quality flags, the remaining coincident analysis time in O1 is 48.3 days with GSTLal analysis, 46.1 with pycbc

        # The second observing run (O2) ran from November 30th, 2016 to August 25th, 2017 --> 267 days
        # During the O2 run the duty cycles were 62% for LIGO Hanford and 61% for LIGO Livingston,
        # so that two detectors were in observing mode 46.4% of the time and at least one detector
        # was in observing mode 75.6% of the time.
        # From https://journals.aps.org/prx/pdf/10.1103/PhysRevX.9.031040 :
        # During O2, the individual LIGO detectors had duty factors of approximately 60% with a LIGO
        # network duty factor of about 45%. Times with significant instrumental disturbances are flagged and removed,
        # resulting in about 118 days of data suitable for coincident analysis

        #self.Tobs= (48.3+118)/365.  # yrs
        self.Tobs= (129+267)/365.  # yrs

    def _get_not_BBHs(self):
        return ['GW170817', ]

    def _name_conditions(self, f ):
        return ( ( 'prior' not in f.split('.')[0] ) &  (f.split('_')[0][:2]=='GW') )

    def _get_name_from_fname(self, fname):
        return fname.split('.')[0].split('_')[0]

    def _load_data_event(self, fname, event, nSamplesUse, which_spins='skip'):

        data_path = os.path.join(fname,  event+'_GWTC-1'+self.post_file_extension)

        with h5py.File(data_path, 'r') as f:

            posterior_samples = f['Overall_posterior']

            m1z = posterior_samples['m1_detector_frame_Msun']
            m2z = posterior_samples['m2_detector_frame_Msun']
            dL = posterior_samples['luminosity_distance_Mpc']
            try:
                w = posterior_samples['weights_bin']
            except Exception as e:
                print(e)
                w = np.ones(1)

            if which_spins=='skip':
                spins=[]
            elif which_spins=='chiEff':
                #print('chi_p not available for O1-O2 data ! ')
                s1 = posterior_samples['spin1']
                s2 = posterior_samples['spin2']
                cost1 = posterior_samples['costilt1']
                cost2 = posterior_samples['costilt2']
                sint1 = np.sqrt(1-cost1**2)
                sint2 = np.sqrt(1-cost2**2)
                chi1z = s1*cost1
                chi2z = s2*cost2
                q = m2z/m1z
                chiEff = (chi1z+q*chi2z)/(1+q)

                chiP = np.max( np.array([s1*sint1, (4*q+3)/(4+3*q)*q*s2*sint2 ]) , axis=0 )

                spins=[chiEff, chiP]
            elif which_spins=='s1s2':
                raise NotImplementedError()
                s1 = posterior_samples['spin1']
                s2 = posterior_samples['spin2']
                spins=[s1,s2]

        # Downsample if needed
        #all_ds = self._downsample( [m1z, m2z, dL, w, *spins,], nSamplesUse)

        #m1z = all_ds[0]
        #m2z= all_ds[1]
        #dL =  all_ds[2]
        #spins = all_ds[4:]
        #ws = all_ds[3]

        return m1z, m2z, dL, spins, w

class O1O2InjectionsData(Data):

    def __init__(self, fname, nInjUse=None,  dist_unit=u.Gpc, ifar_th=1 , which_spins='skip', SNR_th=None ):

        self.dist_unit=dist_unit
        self.m1z, self.m2z, self.dL, self.spins, self.log_weights_sel, self.N_gen, self.Tobs = self._load_data(fname, nInjUse, which_spins=which_spins )
        self.logN_gen = np.log(self.N_gen)
        #self.log_weights_sel = np.log(self.weights_sel)
        assert (self.m1z > 0).all()
        assert (self.m2z > 0).all()
        assert (self.dL > 0).all()
        assert(self.m2z<self.m1z).all()

        #self.Tobs=0.5
        self.chiEff = np.zeros(self.m1z.shape)
        print('Obs time: %s yrs' %self.Tobs )

        self.ifar_th=ifar_th
        #gstlal_ifar, pycbc_ifar, pycbc_bbh_ifar = conditions_arr
        self.condition = np.full(self.m1z.shape, True) #(gstlal_ifar > ifar_th) | (pycbc_ifar > ifar_th) | (pycbc_bbh_ifar > ifar_th)

    def get_theta(self):
        return np.array( [self.m1z, self.m2z, self.dL  ] )

    def _load_data(self, fname, nInjUse, which_spins='skip'):

        with h5py.File(fname, 'r') as f:

            Tobs = (48.3+118)/365. #f.attrs['analysis_time_s']/(365.25*24*3600) # years
            Ndraw = 7.1e07 #f.attrs['total_generated']

            m1 = np.array(f['mass1_source'])
            m2 = np.array(f['mass2_source'])
            z = np.array(f['redshift'])
        #s1z = np.array(f['injections/spin1z'])
        #s2z = np.array(f['injections/spin2z'])
            if which_spins=='skip':
                spins=[]
            elif which_spins=='chiEff':
                chi1z = np.array(f['spin1z'])
                chi2z = np.array(f['spin2z'])
                q = m2/m1
                chiEff = (chi1z+q*chi2z)/(1+q)
                print('chi_p not available for O2 selection effects ! ')
                spins=[chiEff, np.full(chiEff.shape, np.NaN)]
                #raise NotImplementedError()
            elif which_spins=='s1s2':
                raise NotImplementedError()
                s1 = np.array(f['spin1z'])
                s2 = np.array(f['spin2z'])
                spins=[s1,s2]

            p_draw = np.array(f['sampling_pdf'])
            if which_spins=='skip':
                print('Removing factor of 1/2 for each spin dimension from p_draw...')
                p_draw *= 4
            log_p_draw = np.log(p_draw)
            #gstlal_ifar = np.array(f['injections/ifar_gstlal'])
            #pycbc_ifar = np.array(f['injections/ifar_pycbc_full'])
            #pycbc_bbh_ifar = np.array(f['injections/ifar_pycbc_bbh'])

            m1z = m1*(1+z)
            m2z = m2*(1+z)
            dL = np.array(Planck15.luminosity_distance(z).to(self.dist_unit).value)
            #dL = np.array(f['injections/distance']) #in Mpc for GWTC2 !
            #if self.dist_unit==u.Gpc:
            #    dL*=1e-03

            print('Re-weighting p_draw to go to detector frame quantities...')
            myCosmo = Cosmo(dist_unit=self.dist_unit)
            #p_draw/=(1+z)**2
            #p_draw/=myCosmo.ddL_dz(z, Planck15.H0.value, Planck15.Om0, -1., 1., 0) #z, H0, Om, w0, Xi0, n
            log_p_draw -=2*np.log1p(z)
            log_p_draw -= myCosmo.log_ddL_dz(z, Planck15.H0.value, Planck15.Om0, -1., 1., 0. )

            print('Number of total injections: %s' %Ndraw)
            print('Number of injections that pass first threshold: %s' %p_draw.shape[0])

            self.max_z = np.max(z)
            print('Max redshift of injections: %s' %self.max_z)
            return m1z, m2z, dL , spins, log_p_draw , Ndraw, Tobs, #(gstlal_ifar, pycbc_ifar, pycbc_bbh_ifar)

import numpy,import astropy,from astropy