import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import random
import time


class NetworkGenerator:
    def __init__(self, 
                 width=10,
                 height=10,
                 variant="F", 
                 topology="highly_connected",
                 node_drop_fraction=0.1,
                 bottleneck_cluster_count=None,
                 bottleneck_edges_per_link=1):
        
        self.variant = variant.upper() # "F" = Fixed Density, "R" = Random/Custom Density
        self.topology = topology.lower()
        self.width = int(width)
        self.height = int(height)
        self.node_drop_fraction = float(node_drop_fraction)

        # Standard config
        self.node_factor = 0.4 

        # Bottleneck settings
        if bottleneck_cluster_count is None:
            area = self.width * self.height
            self.bottleneck_cluster_count = max(2, int(area / 18))
        else:
            self.bottleneck_cluster_count = int(bottleneck_cluster_count)
            
        self.bottleneck_edges_per_link = int(bottleneck_edges_per_link)

        self.graph = None
        self.active_positions = None 

    def generate(self):
        """Generate a connected network representing rooms in a building."""
        max_attempts = 20 

        for attempt in range(max_attempts):
            self._build_node_mask()
            self._initialize_graph()
            self._add_nodes()

            nodes = list(self.graph.nodes())
            if len(nodes) < 2:
                continue

            # Topology-specific edge construction
            if self.topology == "bottlenecks":
                self._build_bottleneck_clusters(nodes)
            else:
                self._connect_all_nodes_by_nearby_growth(nodes)
                self._add_edges()

            # Cleanup - STRICT MODE
            self._remove_intersections()
            self._enforce_edge_budget()
            
            if not nx.is_connected(self.graph):
                self._force_connect_components()
            
            # Final Safety Check
            self._remove_intersections() 
            
            if nx.is_connected(self.graph):
                return self.graph

        raise RuntimeError("Failed to generate a valid network. Try reducing Void Fraction.")

    # --- INTERNAL HELPERS ---
    def _effective_node_drop_fraction(self):
        base = self.node_drop_fraction
        if self.topology == "highly_connected": return max(0.0, base * 0.8)
        if self.topology == "linear": return min(0.95, base * 1.2)
        return base

    def _build_node_mask(self):
        all_positions = [(x, y) for x in range(self.width + 1) for y in range(self.height + 1)]
        drop_frac = self._effective_node_drop_fraction()
        drop = int(drop_frac * len(all_positions))
        
        deactivated = set(random.sample(all_positions, drop)) if drop > 0 else set()
        self.active_positions = set(all_positions) - deactivated

    def _initialize_graph(self):
        self.graph = nx.Graph()
        # Seed near center
        margin_x = max(1, self.width // 4)
        margin_y = max(1, self.height // 4)
        
        low_x, high_x = margin_x, self.width - margin_x
        low_y, high_y = margin_y, self.height - margin_y
        
        middle_active = [
            (x, y) for (x, y) in self.active_positions 
            if low_x <= x <= high_x and low_y <= y <= high_y
        ]
        
        if middle_active:
            seed = random.choice(middle_active)
        elif self.active_positions:
            seed = random.choice(list(self.active_positions))
        else:
            return 

        self.graph.add_node(tuple(seed))

    def _compute_nodes(self):
        total_possible = (self.width + 1) * (self.height + 1)
        # Variant "F" (Fixed) uses stable logic. Variant "R" (Custom) adds randomization.
        base = self.node_factor if self.variant == "F" else random.uniform(0.3, 0.6)
        scale = {"highly_connected": 1.2, "bottlenecks": 0.85, "linear": 0.75}.get(self.topology, 1.0)
        
        target = int(base * scale * total_possible)
        return min(target, len(self.active_positions))

    def _add_nodes(self):
        total_nodes = self._compute_nodes()
        attempts = 0
        while len(self.graph.nodes()) < total_nodes and attempts < (total_nodes * 20):
            attempts += 1
            x = random.randint(0, self.width)
            y = random.randint(0, self.height)
            
            if (x, y) in self.active_positions and (x, y) not in self.graph:
                self.graph.add_node((x, y))

    def _connect_all_nodes_by_nearby_growth(self, nodes):
        connected = set()
        remaining = set(nodes)
        if not remaining: return
        
        current = random.choice(nodes)
        connected.add(current)
        remaining.remove(current)

        while remaining:
            candidates = []
            for n in remaining:
                for c in connected:
                    if abs(n[0]-c[0]) <= 2 and abs(n[1]-c[1]) <= 2:
                        candidates.append(n)
                        break
            
            if not candidates:
                n = min(remaining, key=lambda r: min(abs(r[0]-c[0]) + abs(r[1]-c[1]) for c in connected))
                candidates.append(n)

            candidate = random.choice(candidates)
            
            neighbors = [c for c in connected if abs(c[0]-candidate[0])<=3 and abs(c[1]-candidate[1])<=3]
            neighbors.sort(key=lambda c: abs(c[0]-candidate[0]) + abs(c[1]-candidate[1]))
            
            n = neighbors[0] if neighbors else random.choice(list(connected))
            
            self.graph.add_edge(n, candidate)
            connected.add(candidate)
            remaining.remove(candidate)

    def _compute_edge_count(self):
        n = len(self.graph.nodes())
        if self.topology == "highly_connected": return int(3.5 * n)
        if self.topology == "bottlenecks": return int(1.8 * n)
        return int(random.uniform(1.2, 2.0) * n)

    def _add_edges(self):
        nodes = list(self.graph.nodes())
        if self.topology == "highly_connected":
            self._add_cluster_dense(nodes, self._compute_edge_count())
        elif self.topology == "linear":
            self._make_linear(nodes)

    def _make_linear(self, nodes):
        nodes_sorted = sorted(nodes, key=lambda x: (x[0], x[1]))
        if not nodes_sorted: return
        prev = nodes_sorted[0]
        for nxt in nodes_sorted[1:]:
            if not self._would_create_intersection(prev, nxt):
                self.graph.add_edge(prev, nxt)
            prev = nxt

    def _add_cluster_dense(self, nodes, max_edges):
        edges_added = 0
        nodes = list(nodes)
        random.shuffle(nodes)
        
        for i in range(len(nodes)):
            for j in range(i + 1, len(nodes)):
                if edges_added >= max_edges: return

                n1, n2 = nodes[i], nodes[j]
                dist = max(abs(n1[0]-n2[0]), abs(n1[1]-n2[1]))
                
                if dist <= 4: 
                    if not self._would_create_intersection(n1, n2):
                        self.graph.add_edge(n1, n2)
                        edges_added += 1

    # --- BOTTLENECK LOGIC ---
    def _build_bottleneck_clusters(self, nodes):
        self.graph.remove_edges_from(list(self.graph.edges()))
        clusters, centers = self._spatial_cluster_nodes(nodes, k=self.bottleneck_cluster_count)
        
        for cluster in clusters:
            if len(cluster) < 2: continue
            self._connect_cluster_by_nearby_growth(cluster)
            self._add_cluster_dense(list(cluster), max_edges=max(1, int(3.5 * len(cluster))))

        order = sorted(range(len(clusters)), key=lambda i: (centers[i][0], centers[i][1]))
        for a_idx, b_idx in zip(order[:-1], order[1:]):
            self._add_bottleneck_links(clusters[a_idx], clusters[b_idx], self.bottleneck_edges_per_link)
        
        if not nx.is_connected(self.graph):
            self._force_connect_components()

    def _force_connect_components(self):
        components = list(nx.connected_components(self.graph))
        while len(components) > 1:
            c1 = list(components[0])
            c2 = list(components[1])
            
            best_pair = None
            min_dist = float('inf')
            
            for u in c1:
                for v in c2:
                    d = (u[0]-v[0])**2 + (u[1]-v[1])**2
                    if d < min_dist:
                        if not self._would_create_intersection(u, v):
                            min_dist = d
                            best_pair = (u, v)
            
            if best_pair:
                self.graph.add_edge(best_pair[0], best_pair[1])
            else:
                pass 
            
            prev_len = len(components)
            components = list(nx.connected_components(self.graph))
            if len(components) == prev_len: break

    def _spatial_cluster_nodes(self, nodes, k):
        nodes = list(nodes)
        if k >= len(nodes): return [[n] for n in nodes], nodes[:]
        centers = random.sample(nodes, k)
        clusters = [[] for _ in range(k)]
        for n in nodes:
            best_i = min(range(k), key=lambda i: max(abs(n[0]-centers[i][0]), abs(n[1]-centers[i][1])))
            clusters[best_i].append(n)
        return clusters, centers

    def _connect_cluster_by_nearby_growth(self, cluster_nodes):
        self._connect_all_nodes_by_nearby_growth(cluster_nodes)

    def _add_bottleneck_links(self, cluster_a, cluster_b, m):
        pairs = []
        for u in cluster_a:
            for v in cluster_b:
                dist = max(abs(u[0]-v[0]), abs(u[1]-v[1]))
                pairs.append((dist, u, v))
        pairs.sort(key=lambda t: t[0])
        
        added = 0
        for _, u, v in pairs:
            if added >= m: break
            if not self.graph.has_edge(u, v) and not self._would_create_intersection(u, v): 
                self.graph.add_edge(u, v)
                added += 1

    # --- GEOMETRY & CLEANUP ---
    def _remove_intersections(self):
        pass_no = 0
        while pass_no < 8:
            pass_no += 1
            edges = list(self.graph.edges())
            intersections = []
            
            for i in range(len(edges)):
                for j in range(i+1, len(edges)):
                    e1 = edges[i]
                    e2 = edges[j]
                    if self._segments_intersect(e1[0], e1[1], e2[0], e2[1]):
                        intersections.append((e1, e2))
            
            if not intersections: break
            
            for e1, e2 in intersections:
                if not self.graph.has_edge(*e1) or not self.graph.has_edge(*e2): continue
                l1 = (e1[0][0]-e1[1][0])**2 + (e1[0][1]-e1[1][1])**2
                l2 = (e2[0][0]-e2[1][0])**2 + (e2[0][1]-e2[1][1])**2
                rem = e1 if l1 > l2 else e2
                self.graph.remove_edge(*rem)

    def _enforce_edge_budget(self):
        budget = self._compute_edge_count()
        while len(self.graph.edges()) > budget:
            edges = list(self.graph.edges())
            rem = random.choice(edges)
            self.graph.remove_edge(*rem)
            if not nx.is_connected(self.graph):
                self.graph.add_edge(*rem)
                break 

    def _segments_intersect(self, a, b, c, d):
        if a == c or a == d or b == c or b == d: return False
        def ccw(A,B,C): return (C[1]-A[1]) * (B[0]-A[0]) > (B[1]-A[1]) * (C[0]-A[0])
        return ccw(a,c,d) != ccw(b,c,d) and ccw(a,b,c) != ccw(a,b,d)

    def _would_create_intersection(self, u, v):
        for a, b in self.graph.edges():
            if u == a or u == b or v == a or v == b: continue
            if self._segments_intersect(u, v, a, b): return True
        return False