/* === File: hnsw/include/hnsw.h === */
#pragma lib "libhnsw.a"
#pragma src "/sys/src/libhnsw"

//hnsw.h

typedef struct Hnsw Hnsw;
typedef struct Hnswgraph Hnswgraph;
typedef struct HnswPQ HnswPQ;



typedef struct HnswSE {
	// node search element (distance to node)
	Hnsw *n;
	double d;
} HnswSE;

typedef struct HnswPQ {
	// priority queue
	HnswSE *l; /* [capacity] */
	//no capacity member, for "scratchpad versions" the cap is efconstruction/ef
		//for node versions it's maxk, when relevant capacity is passed to work functions by graph which knows.
		//we simply can't waste the amount of memory having them here would 
	uint size;
}HnswPQ;

typedef struct Hnsw {
	//user node annonomous embed
	u64int s; //search id
	HnswPQ * n; //HnswPQ neighbors[layer+1] ; each with maxk capacity
	uint layer; 
}Hnsw;

typedef struct HnswLayer {
	Hnsw * n;
	uint size; //number currently added 
	uint capacity; 
} HnswLayer ;

typedef struct HnswNodeArena HnswNodeArena;
struct HnswNodeArena{
	Hnsw ** nodes;
	HnswNodeArena * next;
} ;

typedef struct Hnswgraph {
	/* creation parameters */
	
	/* provide your own magnitude distance function for your nodes.  */
	double (*dist)(Hnsw*, Hnsw*);
	double ml; 
	int maxk;
	int efcons;
	int created;

	u64int cs; // current search id
	int topl;
	
	//state
	Hnsw	*ep;
	
	HnswNodeArena * arenahead;
	HnswNodeArena * nodearena;
	Hnsw	**allnodes; // malloced based on estimated max number of nodes at creation time
	uint nodecount;
	uint arenaindex;
	uint firstarenasize;
	uint growstep; // when reaching max nodes the number of node entries to try to grow to. 


/* neighbor culling is asymetrical, the 'dangling' links are definitionally weak and likely to be into dense clusters from outside them. search would likely only take that one way link if the direction it wanted to go is "into" that cluster anyway, so this is fine. */

}Hnswgraph;

//example values (50000 , 10000, 50, 100, .6, L2dist)
Hnswgraph *hnswcreate(int estmaxnodes, int growstep, int maxk, int efcons, double ml, double (*dist)(Hnsw*, Hnsw*));
int hnswinsert(Hnswgraph * g, Hnsw * n);
void hnswfreepq(HnswPQ * pq);
HnswPQ * hnswsearch(Hnswgraph *g, Hnsw * n, uint ef);

/* === File: hnsw/src/libhnsw/hnsw.c === */
/* 
	see : raw.githubusercontent.com/brtholomy/hnsw/refs/heads/master/hnsw.py 
		god the python is such shit.
we echew a lot of the complexity by assuming nothing about the data 
	( in the same way as avl.h )
*/
//hnsw.c
#include <u.h>
#include <libc.h>
#include <hnsw.h>

#define HNSWCREATED 0x484e5357
 
void
removehnswpqnode(HnswPQ	* q, uint i) {
	if (!q || !q->l)
		return; //sanity check
	if ( i < q->size ) { 
		q->size = q->size - 1;
		if (i == (q->size))
			return;
		memmove( &(q->l[i]), &(q->l[i+1]), sizeof(HnswSE) * (q->size - i));
	}
	return;
}

/*
	maintains strict ordering, if full and n is closer than furtheest culls the "last" node to make room
*/
void
inserthnswpqnode(HnswPQ * q,  HnswSE * n, int cap, int hightolow ){
	if(!q || !q->l)
		return;
	int loc = -1;
		//no matter what we need to find where we may add the item
	for(int i = 0 ; loc < 0 && i < cap && i < q->size; ++i){
		if (hightolow){
			loc = (q->l[i].d < n->d) ? i : loc;
		} else { 
			loc = (q->l[i].d > n->d) ? i : loc;
		}
	}

	if ( loc >= 0 ) {
		if (q->size < cap){
			memmove(&(q->l[loc+1]), &(q->l[loc]), sizeof(HnswSE) * (q->size - loc));
			q->size = q->size + 1;
		} else {
			memmove(&(q->l[loc+1]), &(q->l[loc]), sizeof(HnswSE) * (cap - 1 - loc));
		}
		q->l[loc] = *n; 
		return;
	}
	if( q->size < cap){
		q->l[q->size] = *n;
		q->size = q->size + 1;
	}

} 

void
hnswallocatenodequeue(Hnsw * n, uint c){
	uint sizePQxL   = (n->layer+1) * sizeof(HnswPQ);
	uint sizeSExC   = c * sizeof(HnswSE);
	uint sizeSExCxL = (n->layer+1) * sizeSExC;
	n->n = malloc(sizePQxL + sizeSExCxL);
	u8int * endofpqs = ((u8int *)n->n)+sizePQxL;
	for(int i=0 ; i <= n->layer ; ++i){
		n->n[i].l = (HnswSE *) (endofpqs + i*sizeSExC); 
		n->n[i].size = 0;
	}
}

//node needs to have it's neighbors set this adds node to each neighbor
// culling as needed - (will remove furthest node from neighbor if needed, and replace with n)
// if all nodes in neighbor list are closer than n then that neighbor is culled from n instead. 
// "what if all neighbors are culled from n?"  - set maxk higher next time. (skill issue, apparently)
void
addedges(Hnswgraph *g, Hnsw *n, uint layer){
//	print("inserting edges on layer(%d)\n", layer);
	if(!n || n->layer < layer){
		return; //bru
	}
	HnswPQ * nl = &(n->n[layer]);
	for(int i = nl->size-1; i >=0; --i){
//		print("addedge- i(%d) - nl->size(%d)\n",i, nl->size);
	//loop backward so we can pop neighbors without fucking up our loop
		HnswSE * nse = &(nl->l[i]);
		if(nse->n == nil)
			continue; // just skip it. this shouldn't happen (TODO FIND ROOT CAUSE)
//		print("addedge - nse->n(%p)\n",nse->n);
//		print("addedge - nse->n->layers(%d)\n", nse->n->layer);
		HnswPQ * mpq = &(nse->n->n[layer]);
		HnswSE ntom = {.n=n,.d=nse->d};

		if (mpq->size < g->maxk){
		//just add it to neighbors list
			
			inserthnswpqnode(mpq, &ntom, g->maxk, 0);
			continue;
		}
		if (mpq->l[mpq->size-1].d < nse->d){
		//neighbor's furthests is closer. git fucked nse
			removehnswpqnode(nl, i);
			continue;
		}
		//neighborlist is full and n is closer than furthest
		inserthnswpqnode( mpq, &ntom, g->maxk, 0);
	}
}

void
_hnswsearchlayer (Hnswgraph *g, HnswPQ *nearestn, HnswPQ *candidates, Hnsw * n, HnswSE * ep, uint layer, uint ef){

	inserthnswpqnode(candidates, ep, 8*ef, 0);
	inserthnswpqnode(nearestn, ep, ef, 0);

	while ( candidates->size > 0 ){
		HnswSE c = candidates->l[0];
		removehnswpqnode(candidates, 0);
		HnswSE *f = &(nearestn->l[nearestn->size - 1]);
		if ( c.d > f->d ) //all remaining candidates too far. done
			break;
//		print("search size(%d) - c(%p) - ", candidates->size, c);
//		print("c.n(%p) - ", c.n);
//		print("c.n->layer(%d)\n", c.n->layer);
		HnswPQ * cn = &(c.n->n[layer]); //grab this candidates neighboors on this layer
//		print("cn(%p)\n",cn);
		for(int i=0; i < (cn->size) ; ++i){
//			print("i(%d)\n",i);
			if (cn->l[i].n->s < g->cs) {
				HnswSE curn = {.n=cn->l[i].n, .d= g->dist(n, cn->l[i].n)};
				curn.n->s = g->cs; //visit each
				f = &(nearestn->l[nearestn->size - 1]);
				if ( curn.d < f->d || nearestn->size < ef){
					inserthnswpqnode(candidates, &curn, 8*ef, 0);
					inserthnswpqnode(nearestn, &curn, ef, 0);
				}
			}
		}
		
		
	}
	return;
}

HnswPQ *
hnswsearch(Hnswgraph *g, Hnsw * n, uint ef){
	HnswPQ	*nearestn = malloc(sizeof(HnswPQ));
	nearestn->l = malloc(ef * sizeof(HnswSE));
	nearestn->size = 0;
	HnswPQ	candidates = {.l = malloc(8 * ef * sizeof(HnswSE)), .size = 0 };

	HnswSE ep = {.n=g->ep, .d=g->dist(n, g->ep)};
	g->cs = g->cs + 1; //search scope started for insert
	ep.n->s = g->cs; //distance already calculated so ep has been "visited"

	for (int curlayer = g->topl; curlayer >= 0 ; --curlayer){
		candidates.size=0;
		nearestn->size=0;
		_hnswsearchlayer(g, nearestn, &candidates, n, &ep, curlayer, ef);
		ep = nearestn->l[0];
	}
	free (candidates.l);
	return nearestn;
}

/* use this to free the PQ returned by search */
void
hnswfreepq(HnswPQ * pq){
	free(pq->l);
	free(pq);
}


void
hnswaddtoarena(Hnswgraph * g, Hnsw * n){
	if (g->nodecount < g->firstarenasize){
		g->nodearena->nodes[g->nodecount++] = n;
		if(g->nodecount == g->firstarenasize){
			g->nodearena->next = malloc(sizeof(HnswNodeArena));
			g->nodearena->next->nodes = malloc(sizeof(Hnsw*)*(g->growstep));
			g->nodearena->next->next = nil;
			g->nodearena = g->nodearena->next;
			g->arenaindex = 0;
		}
		return;
	}
	HnswNodeArena * arena = g->nodearena;
	arena->nodes[ g->arenaindex] = n;
	++(g->nodecount);
	++(g->arenaindex);
	if( g->arenaindex == g->growstep-1){
		arena->next = malloc(sizeof(HnswNodeArena));
		arena->next->nodes = malloc(sizeof(Hnsw*)*(g->growstep));
		arena->next->next = nil;
		g->nodearena = g->nodearena->next;
		g->arenaindex = 0;
	}
}

//inserting is *not* threadsafe. neither with searchs nor with other inserts.
int
hnswinsert(Hnswgraph * g, Hnsw * n) {
	u8int newep = 0; //flag to update graph ep if this is the new one
	if(!g || g->created != HNSWCREATED || !n)
		return -1;
	HnswSE epc = {.n=g->ep, .d=g->ep?g->dist(n, g->ep):0.0};
	n->s = 0;
	n->layer = (uint) floor(-1 * log(frand()) * g->ml );
//	print("add node @layer(%d) - topl(%d)\n", n->layer, g->topl );
	if (((int)n->layer) > g->topl) { 
		n->layer = g->topl = g->topl + 1;
//		print("new toplayer - %d \n", g->topl);
		newep = 1; //technically we can set it here, but it's not a fully graphed node yet. better wait till it is. 
	}
	hnswallocatenodequeue(n, g->maxk);
	hnswaddtoarena(g, n);
	if(epc.n == nil){ //this is literally the first node, we only need set ep now
//		print("first node( %d)\n", g->topl);
		g->ep = n;
		return 0;
	}



	HnswPQ	candidates = {.l = malloc(8 * g->efcons * sizeof(HnswSE)), .size = 0 };
	HnswPQ	nearestn = {.l = malloc(g->efcons * sizeof(HnswSE)), .size = 0 };

	g->cs = g->cs + 1; //search scope started for insert
	epc.n->s = g->cs; //distance already calculated so ep has been "visited"
	//from the top to layer, search to find the ep on the node's layer nearest the node.
	for (int i = (g->topl - newep); i > n->layer ; --i){
//		print("insert layer ( %d)\n", i);                                /////////////////////////////////
		_hnswsearchlayer(g, &nearestn, &candidates, n, &epc, i, 1);
		epc = nearestn.l[0];
		candidates.size=0;
		nearestn.size=0;
	}//this will not happen if the node is the new highest node. This is fine. 
	
	for (int i = (n->layer - newep)  ; i >= 0 ; --i){
		_hnswsearchlayer(g, &nearestn, &candidates, n, &epc, i, g->efcons);
		for(int j = 0; j < g->maxk && j < g->efcons && j < nearestn.size ; ++j){
			n->n[i].l[j]=nearestn.l[j];
			n->n[i].size=j+1;
		}
		addedges(g, n, i); //handles shrinking edges when needed
		epc = n->n[i].l[0];
		candidates.size=0;
		nearestn.size=0;
	}

	free(candidates.l);
	free(nearestn.l);

	if(newep)
		g->ep = n;
	return 0;
	
}



Hnswgraph *
hnswcreate(int estmaxnodes, int growstep, int maxk, int efcons, double ml, double (*dist)(Hnsw*, Hnsw*)){
	Hnswgraph * g = malloc(sizeof(Hnswgraph));
	{
		Hnswgraph t = 
			{.dist=dist, .ml=ml, .maxk=maxk, .efcons=efcons, .cs=1, .topl=-1, .ep=nil, .nodecount=0, .firstarenasize=estmaxnodes, .growstep=growstep};
		*g = t;
	}
//	g->allnodes = malloc(sizeof(Hnsw*)*(estmaxnodes+1)); //extra to linked list allocate in steps.  
	g->nodearena = malloc(sizeof(HnswNodeArena));
	g->nodearena->nodes = malloc(sizeof(Hnsw*)*(estmaxnodes));
	g->nodearena->next = nil;
	g->arenahead = g->nodearena;
	g->created=HNSWCREATED;
	return g; 
}