"""utilities holding temporal network info"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/tnet.ipynb.

# %% auto 0
__all__ = ['TempNet', 'nimage_from_npair']

# %% ../nbs/API/tnet.ipynb 3
import math
import numpy as np
import zarr
from .utils_ import ngjit

# %% ../nbs/API/tnet.ipynb 4
def _imagepair_from_bandwidth(nimages,bandwidth=None):
    if bandwidth is None: bandwidth = nimages
    assert nimages >= bandwidth
    ref, sec = np.triu_indices(nimages, 1)
    idx = np.where((sec-ref)<=bandwidth)
    return np.stack((ref[idx], sec[idx]),axis=-1).astype(np.int32)

# %% ../nbs/API/tnet.ipynb 5
@ngjit
def _are_edges_sorted(edges):
    num_edges = edges.shape[0]
    for i in range(num_edges - 1):
        if edges[i, 0] > edges[i + 1, 0] or (edges[i, 0] == edges[i + 1, 0] and edges[i, 1] > edges[i + 1, 1]):
            return False
    return True

# %% ../nbs/API/tnet.ipynb 6
@ngjit
def _are_edges_directed(edges):
    '''directed = reference index smaller than secondary index'''
    num_edges = edges.shape[0]
    for i in range(num_edges):
        if edges[i, 0] >= edges[i, 1]:
            return False
    return True

# %% ../nbs/API/tnet.ipynb 7
@ngjit
def _are_edges_connected(edges):
    num_edges = edges.shape[0]

    current_ref = -1
    for i in range(num_edges):
        if edges[i,0] == current_ref:
            continue
        else:
            if edges[i,0] != edges[i,1]-1:
                return False
            current_ref = edges[i,0]
    return True

# %% ../nbs/API/tnet.ipynb 8
class TempNet(object):
    def __init__(self,image_pairs,check_if_valid=True):
        if check_if_valid:
            are_edges_sorted = _are_edges_sorted(image_pairs)
            if not are_edges_sorted:
                raise ValueError('input image pairs are not sorted.')
            are_edges_directed = _are_edges_directed(image_pairs)
            if not are_edges_directed:
                raise ValueError('input image pairs are not directed (reference index larger than or equal to secondary index).')
            are_edges_connected = _are_edges_connected(image_pairs)
            if not are_edges_connected:
                raise ValueError('input image pairs are not connected.')
        self._image_pairs = image_pairs.astype(np.int32)
        self._n_images = self._image_pairs[-1,-1]+1
        self._image_pairs_1d = self._image_pairs[:,0]*self._n_images+self._image_pairs[:,1]

    @classmethod
    def from_bandwidth(cls, nimages, bandwidth=None):
        image_pairs = _imagepair_from_bandwidth(nimages,bandwidth)
        return cls(image_pairs,check_if_valid=False)

    @property
    def image_pairs(self):
        return self._image_pairs

    def image_pairs_idx(
        self,
        image_pairs=None,
        ref=None,
        sec=None,
    ):
        single_input = False
        if image_pairs is not None:
            ref = image_pairs[:,0]
            sec = image_pairs[:,1]
        else:
            if isinstance(ref,int):
                assert isinstance(sec,int)
                single_input = True
                ref = np.array([ref,]); sec = np.array([sec,])
            else:
                ref = np.array(ref); sec = np.array(sec)
            assert len(ref) == len(sec)
        
        query_1d = ref*self._n_images+sec
        mask = np.isin(query_1d,self._image_pairs_1d,assume_unique=True,invert=True)
        result = np.searchsorted(self._image_pairs_1d,query_1d)
        result[mask] = -1
        if single_input:
            result = result[0]
        return result

    def save(self,path:str, # zarr path
            ):
        '''Save the TempNet.'''
        tempnet_zarr = zarr.open(path,mode='w',shape=self._image_pairs.shape,dtype=self._image_pairs.dtype)
        tempnet_zarr[:] = self._image_pairs[:]

    @classmethod
    def load(cls, zarr_path:str, # zarr path
            ):
        '''classmethod to load the saved HilbertRtree.'''
        tempnet_zarr = zarr.open(zarr_path,mode='r')
        return cls(tempnet_zarr[:], check_if_valid=False)

# %% ../nbs/API/tnet.ipynb 17
def nimage_from_npair(npair):
    '''infer number of images from number of pairs assuming all pairs are selected.'''
    possible_nimage = math.ceil(math.sqrt(2*npair))
    for i in [0,-1,1,-2,2]:
        possible_nimage = possible_nimage+i
        if possible_nimage*(possible_nimage-1) == 2*npair:
            return possible_nimage
    raise ValueError('input is not correct value of number of all pairs')
