graph.go

package mlgraph

import (
	"fmt"
	"sync"
)

// Edge represents a connection between nodes with metadata
type Edge[K comparable] struct {
	From   K
	To     K
	Weight float64
	Data   map[string]any
}

// Graph represents a directed graph with nodes and edges
type Graph[K comparable, V any] struct {
	mu    sync.RWMutex
	nodes map[K]V
	edges map[K][]Edge[K]
}

// New creates a new graph instance
func New[K comparable, V any]() *Graph[K, V] {
	return &Graph[K, V]{
		nodes: make(map[K]V),
		edges: make(map[K][]Edge[K]),
	}
}

// AddNode adds a node to the graph
func (g *Graph[K, V]) AddNode(id K, value V) error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.nodes[id]; exists {
		return fmt.Errorf("node %v already exists", id)
	}

	g.nodes[id] = value
	g.edges[id] = []Edge[K]{}
	return nil
}

// UpdateNode updates an existing node's value
func (g *Graph[K, V]) UpdateNode(id K, value V) error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.nodes[id]; !exists {
		return fmt.Errorf("node %v does not exist", id)
	}

	g.nodes[id] = value
	return nil
}

// GetNode retrieves a node's value
func (g *Graph[K, V]) GetNode(id K) (V, bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	value, exists := g.nodes[id]
	return value, exists
}

// RemoveNode removes a node and all its edges
func (g *Graph[K, V]) RemoveNode(id K) error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.nodes[id]; !exists {
		return fmt.Errorf("node %v does not exist", id)
	}

	// Remove the node
	delete(g.nodes, id)
	delete(g.edges, id)

	// Remove edges pointing to this node
	for nodeID, edges := range g.edges {
		filtered := []Edge[K]{}
		for _, edge := range edges {
			if edge.To != id {
				filtered = append(filtered, edge)
			}
		}
		g.edges[nodeID] = filtered
	}

	return nil
}

// AddEdge adds a directed edge between two nodes
func (g *Graph[K, V]) AddEdge(from, to K, weight float64) error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.nodes[from]; !exists {
		return fmt.Errorf("source node %v does not exist", from)
	}
	if _, exists := g.nodes[to]; !exists {
		return fmt.Errorf("target node %v does not exist", to)
	}

	edge := Edge[K]{
		From:   from,
		To:     to,
		Weight: weight,
		Data:   make(map[string]any),
	}

	g.edges[from] = append(g.edges[from], edge)
	return nil
}

// AddEdgeWithData adds a directed edge with metadata
func (g *Graph[K, V]) AddEdgeWithData(from, to K, weight float64, data map[string]any) error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.nodes[from]; !exists {
		return fmt.Errorf("source node %v does not exist", from)
	}
	if _, exists := g.nodes[to]; !exists {
		return fmt.Errorf("target node %v does not exist", to)
	}

	edge := Edge[K]{
		From:   from,
		To:     to,
		Weight: weight,
		Data:   data,
	}

	g.edges[from] = append(g.edges[from], edge)
	return nil
}

// GetEdges returns all edges from a node
func (g *Graph[K, V]) GetEdges(from K) ([]Edge[K], error) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if _, exists := g.nodes[from]; !exists {
		return nil, fmt.Errorf("node %v does not exist", from)
	}

	edges := make([]Edge[K], len(g.edges[from]))
	copy(edges, g.edges[from])
	return edges, nil
}

// RemoveEdge removes an edge between two nodes
func (g *Graph[K, V]) RemoveEdge(from, to K) error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.nodes[from]; !exists {
		return fmt.Errorf("source node %v does not exist", from)
	}

	filtered := []Edge[K]{}
	found := false
	for _, edge := range g.edges[from] {
		if edge.To != to {
			filtered = append(filtered, edge)
		} else {
			found = true
		}
	}

	if !found {
		return fmt.Errorf("edge from %v to %v does not exist", from, to)
	}

	g.edges[from] = filtered
	return nil
}

// HasEdge checks if an edge exists between two nodes
func (g *Graph[K, V]) HasEdge(from, to K) bool {
	g.mu.RLock()
	defer g.mu.RUnlock()

	edges, exists := g.edges[from]
	if !exists {
		return false
	}

	for _, edge := range edges {
		if edge.To == to {
			return true
		}
	}
	return false
}

// Nodes returns all node IDs in the graph
func (g *Graph[K, V]) Nodes() []K {
	g.mu.RLock()
	defer g.mu.RUnlock()

	nodes := make([]K, 0, len(g.nodes))
	for id := range g.nodes {
		nodes = append(nodes, id)
	}
	return nodes
}

// NodeCount returns the number of nodes
func (g *Graph[K, V]) NodeCount() int {
	g.mu.RLock()
	defer g.mu.RUnlock()
	return len(g.nodes)
}

// EdgeCount returns the total number of edges
func (g *Graph[K, V]) EdgeCount() int {
	g.mu.RLock()
	defer g.mu.RUnlock()

	count := 0
	for _, edges := range g.edges {
		count += len(edges)
	}
	return count
}

// BFS performs a breadth-first search from the start node
func (g *Graph[K, V]) BFS(start K, visit func(K, V) bool) error {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if _, exists := g.nodes[start]; !exists {
		return fmt.Errorf("start node %v does not exist", start)
	}

	visited := make(map[K]bool)
	queue := []K{start}

	for len(queue) > 0 {
		current := queue[0]
		queue = queue[1:]

		if visited[current] {
			continue
		}
		visited[current] = true

		value := g.nodes[current]
		if !visit(current, value) {
			break
		}

		for _, edge := range g.edges[current] {
			if !visited[edge.To] {
				queue = append(queue, edge.To)
			}
		}
	}

	return nil
}

// DFS performs a depth-first search from the start node
func (g *Graph[K, V]) DFS(start K, visit func(K, V) bool) error {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if _, exists := g.nodes[start]; !exists {
		return fmt.Errorf("start node %v does not exist", start)
	}

	visited := make(map[K]bool)
	var dfsHelper func(K) bool

	dfsHelper = func(current K) bool {
		if visited[current] {
			return true
		}
		visited[current] = true

		value := g.nodes[current]
		if !visit(current, value) {
			return false
		}

		for _, edge := range g.edges[current] {
			if !dfsHelper(edge.To) {
				return false
			}
		}
		return true
	}

	dfsHelper(start)
	return nil
}

// TopologicalSort returns nodes in topological order (for DAGs)
func (g *Graph[K, V]) TopologicalSort() ([]K, error) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	// Calculate in-degrees
	inDegree := make(map[K]int)
	for node := range g.nodes {
		if _, exists := inDegree[node]; !exists {
			inDegree[node] = 0
		}
		for _, edge := range g.edges[node] {
			inDegree[edge.To]++
		}
	}

	// Find nodes with no incoming edges
	queue := []K{}
	for node, degree := range inDegree {
		if degree == 0 {
			queue = append(queue, node)
		}
	}

	result := []K{}
	for len(queue) > 0 {
		current := queue[0]
		queue = queue[1:]
		result = append(result, current)

		for _, edge := range g.edges[current] {
			inDegree[edge.To]--
			if inDegree[edge.To] == 0 {
				queue = append(queue, edge.To)
			}
		}
	}

	if len(result) != len(g.nodes) {
		return nil, fmt.Errorf("graph contains a cycle")
	}

	return result, nil
}

// HasCycle detects if the graph contains a cycle
func (g *Graph[K, V]) HasCycle() bool {
	g.mu.RLock()
	defer g.mu.RUnlock()

	visited := make(map[K]int) // 0: unvisited, 1: visiting, 2: visited

	var hasCycle func(node K) bool
	hasCycle = func(node K) bool {
		visited[node] = 1 // Mark as visiting

		for _, edge := range g.edges[node] {
			if visited[edge.To] == 1 {
				return true // Back edge found
			}
			if visited[edge.To] == 0 && hasCycle(edge.To) {
				return true
			}
		}

		visited[node] = 2 // Mark as visited
		return false
	}

	for node := range g.nodes {
		if visited[node] == 0 {
			if hasCycle(node) {
				return true
			}
		}
	}

	return false
}

// Clone creates a deep copy of the graph
func (g *Graph[K, V]) Clone() *Graph[K, V] {
	g.mu.RLock()
	defer g.mu.RUnlock()

	newGraph := New[K, V]()

	// Copy nodes
	for id, value := range g.nodes {
		newGraph.nodes[id] = value
	}

	// Copy edges
	for from, edges := range g.edges {
		newEdges := make([]Edge[K], len(edges))
		for i, edge := range edges {
			newData := make(map[string]any)
			for k, v := range edge.Data {
				newData[k] = v
			}
			newEdges[i] = Edge[K]{
				From:   edge.From,
				To:     edge.To,
				Weight: edge.Weight,
				Data:   newData,
			}
		}
		newGraph.edges[from] = newEdges
	}

	return newGraph
}