export class GraphEdge<T, V> {
    public source: GraphNode<T, V>
    public target: GraphNode<T, V>
    public weight?: number
    public context?: V

    constructor(source: GraphNode<T, V>, target: GraphNode<T, V>, weight?: number, context?: V) {
        this.source = source
        this.target = target
        this.weight = weight ?? 1
        this.context = context
    }

    public get reverse(): GraphEdge<T, V> {
        return new GraphEdge(this.target, this.source, this.weight)
    }
}

export class GraphNode<T, V> {
    public id: number
    public edges: Set<GraphEdge<T, V>>
    public context?: T

    constructor(id: number, context?: T) {
        this.id = id
        this.edges = new Set()
        this.context = context
    }

    public get degree(): number {
        const hashEdge = (sourceId: number, targetId: number): string => {
            const [minId, maxId] = sourceId < targetId ? [sourceId, targetId] : [targetId, sourceId]
            return `${minId}-${maxId}`
        }

        const edgeHashes = new Set<string>()
        for (const edge of this.edges) {
            edgeHashes.add(hashEdge(this.id, edge.target.id))
        }

        return edgeHashes.size

    }

    public get neighbors(): GraphNode<T, V>[] {
        return Array.from(this.edges).map(edge => edge.target)
    }

    public addEdge(target: GraphNode<T, V>, weight: number = 1, context?: V): void {
        const edge1 = new GraphEdge(this, target, weight, context)
        const edge2 = new GraphEdge(target, this, weight, context) // Add edge in both directions for undirected graph
        this.edges.add(edge1)
        target.edges.add(edge2)
    }

    public removeEdge(target: GraphNode<T, V>): void {
        this.edges = new Set([...this.edges].filter(edge => edge.target !== target))
    }

    public hasEdge(target: GraphNode<T, V>): boolean {
        return [...this.edges].some(edge => edge.target === target)
    }

    public getEdge(target: GraphNode<T, V>): GraphEdge<T, V> | undefined {
        return [...this.edges].find(edge => edge.target === target)
    }
}


export default class Graph<T, V> {
    public nodes: Map<number, GraphNode<T, V>> = new Map()

    public static fromNodes<T, V>(nodes: GraphNode<T, V>[]): Graph<T, V> {
        const graph = new Graph<T, V>()
        for (const node of nodes) {
            graph.nodes.set(node.id, node)
        }
        return graph
    }

    public *[Symbol.iterator]() {
        for (const node of this.nodes) {
            yield node
        }
    }

    public forEachNode(callback: (node: GraphNode<T, V>, i: number) => void): void {
        let i = 0
        for (const node of this.nodes.values()) {
            callback(node, i)
            i++
        }
    }

    public forEachEdge(callback: (edge: GraphEdge<T, V>, i: number) => void): void {
        let i = 0
        for (const node of this.nodes.values()) {
            for (const edge of node.edges) {
                callback(edge, i)
                i++
            }
        }
    }

    public nodeCount(): number {
        return this.nodes.size
    }

    public edgeCount(): number {
        return Array.from(this.nodes.values()).reduce((acc, node) => acc + node.degree, 0)
    }

    public asArray(): GraphNode<T, V>[] {
        return Array.from(this.nodes.values())
    }

    public createNode(id: number, context?: T): GraphNode<T, V> {
        if (this.nodes.has(id)) {
            return this.nodes.get(id)!
        }

        const node = new GraphNode<T, V>(id, context)
        this.nodes.set(id, node)
        return node
    }

    public addNode(node: GraphNode<T, V>): GraphNode<T, V> {
        this.nodes.set(node.id, node)
        return node
    }

    public getNode(id: number): GraphNode<T, V> | undefined {
        return this.nodes.get(id)
    }

    public hasNode(id: number): boolean {
        return this.nodes.has(id)
    }

    public createEdge(source: GraphNode<T, V>, target: GraphNode<T, V>, weight: number = 1, context?: V): void {
        source.addEdge(target, weight, context)
        target.addEdge(source, weight, context) // Add edge in both directions for undirected graph
    }

    public removeEdge(source: GraphNode<T, V>, target: GraphNode<T, V>): void {
        source.removeEdge(target)
        target.removeEdge(source) // Remove edge in both directions for undirected graph
    }

    public hasEdge(source: GraphNode<T, V>, target: GraphNode<T, V>): boolean {
        return source.hasEdge(target)
    }

    public getEdge(source: GraphNode<T, V>, target: GraphNode<T, V>): GraphEdge<T, V> | undefined {
        return source.getEdge(target)
    }

    public getEdgesFromNode(node: GraphNode<T, V>): GraphEdge<T, V>[] {
        return [...node.edges].filter(edge => edge.source === node)
    }

    public getEdgesToNode(node: GraphNode<T, V>): GraphEdge<T, V>[] {
        return [...node.edges].filter(edge => edge.target === node)
    }

    public adjacencyMatrix(): number[][] {
        const nodeIds = Array.from(this.nodes.keys())
        const numNodes = nodeIds.length
        const matrix: number[][] = Array(numNodes).fill(null).map(() => Array(numNodes).fill(0))
    
        for (let i = 0; i < numNodes; i++) {
            const sourceId = nodeIds[i]
            const source = this.getNode(sourceId)
    
            if (source) {
                for (let j = 0; j < numNodes; j++) {
                    const targetId = nodeIds[j]
                    if (source.hasEdge(this.getNode(targetId))) {
                        const edge = source.getEdge(this.getNode(targetId))
                        matrix[i][j] = edge ? edge.weight : 1
                    } else {
                        matrix[i][j] = 0
                    }
                }
            }
        }
    
        return matrix
    }

    /**
     * Gives the connectivity metrics of the graph.
     * 
     * **Alpha**: The higher the alpha index, the more a network is connected. Trees and simple networks will have a value of 0. A value of 1 indicates a completely connected network.
     * 
     * **Beta**: Trees and simple networks have Beta value of less than one. A connected network with one cycle has a value of 1.
     * 
     * **Gamma**: Considers the relationship between the number of observed links and the number of possible links. The value of gamma is between 0 and 1 where a value of 1 indicates a completely connected network.
     */
    public getConnectivityMetrics(): { alpha: number, beta: number, gamma: number } {
        const networkEdges = this.edgeCount()
        const networkNodes = this.nodeCount()

        const alpha = (networkEdges - networkNodes + 1) / (2 * networkNodes - 5)
        const beta = networkEdges / networkNodes
        const gamma = networkEdges / (3 * (networkNodes - 2))

        return { alpha, beta, gamma }
    }

    /**
     * Returns the shortest path between two nodes in the graph using Dijkstra's algorithm.
     */
    public shortestPath(
        source: GraphNode<T, V>,
        target: GraphNode<T, V>,
        metric: (edge: GraphEdge<T, V>) => number = edge => edge.weight
    ): GraphNode<T, V>[] {
        const unvisited = new Set(this.nodes.values())
        const distances = new Map<GraphNode<T, V>, number>()
        const previous = new Map<GraphNode<T, V>, GraphNode<T, V>>()
    
        distances.set(source, 0)
    
        while (unvisited.size > 0) {
            // Find the node with the smallest distance
            let current: GraphNode<T, V> | undefined = undefined
            let smallestDistance = Infinity
            for (const node of unvisited) {
                const distance = distances.get(node) ?? Infinity
                if (distance < smallestDistance) {
                    smallestDistance = distance
                    current = node
                }
            }
    
            // If the smallest distance is Infinity, we are stuck
            if (current === undefined || smallestDistance === Infinity) {
                break
            }
    
            unvisited.delete(current)
    
            // If we reached the target, reconstruct the path
            if (current === target) {
                const path: GraphNode<T, V>[] = []
                while (current !== undefined) {
                    path.unshift(current)
                    current = previous.get(current)
                }
                return path
            }
    
            // Update distances for neighbors
            for (const neighbor of current.neighbors) {
                if (!unvisited.has(neighbor)) {
                    continue
                }
    
                const edge = current.getEdge(neighbor)
                if (!edge) {
                    continue
                }
    
                const distance = distances.get(current)! + metric(edge)
                if (distance < (distances.get(neighbor) ?? Infinity)) {
                    distances.set(neighbor, distance)
                    previous.set(neighbor, current)
                }
            }
        }
    
        // If we exit the loop without finding the target, return an empty path
        return []
    }
}