backend/app/services/graph_service.py

"""
Engine 2 — Supply Chain Knowledge Graph with PageRank + Louvain Community Detection.

Maintains a directed graph of drugs, APIs, rare earths, and production provinces.
Provides:
  - Personalized PageRank for shock propagation (replaces naive BFS)
  - Louvain community detection for supply cluster identification
  - HHI concentration scoring
  - Combined risk formula: R_i = PR_i × (1 - S_i) × e^(-B_i/τ) × C_community

Implements:
  - GNN-based shock propagation concepts from AAAI 2025
  - Causal ML supply chain intervention from arXiv 2024
"""

import math
import logging
from typing import List, Dict, Optional, Any, Literal, Tuple

import networkx as nx

logger = logging.getLogger("shockmap.graph")

# ─── Sector Lead-Time Constants (τ in days) ──────────────────────────────────
SECTOR_TAU = {
    "pharma": 30,
    "rare_earth": 60,
    "semiconductor": 45,
    "ev_battery": 50,
    "solar": 40,
    "defence": 90,
}


class GraphService:
    """
    Engine 2: Maintains a directed graph representation of the multi-sector
    supply chain with PageRank propagation and Louvain community detection.
    """

    def __init__(self, data_loader: Any):
        """
        Initializes the GraphService and builds the DiGraph from seed data.

        Args:
            data_loader: The data provider for drugs, APIs, rare earths, and dependencies.
        """
        self.data_loader = data_loader
        self.graph = nx.DiGraph()
        self._undirected = None  # For Louvain (needs undirected graph)
        self._communities = {}   # node_id → community_id
        self._community_labels = {}  # community_id → descriptive label
        self._community_detection_enabled = True
        self._community_detection_method = "louvain"
        self._pagerank_cache = {}
        self._build_graph()
        self._detect_communities()

    # ─── Graph Construction ──────────────────────────────────────────────

    def _build_graph(self):
        """Constructs the NetworkX graph from data_loader seed data."""
        # Add Drug nodes
        for drug in self.data_loader.get_drugs():
            self.graph.add_node(
                drug.id,
                type="drug",
                name=drug.name,
                sector="pharma",
                nlem_tier=drug.nlem_tier,
                patient_population_estimate=drug.patient_population_estimate,
                has_substitute=drug.has_substitute,
                therapeutic_class=drug.therapeutic_class,
                current_risk=drug.current_risk or 0.0,
                # Engine 2 enrichments
                substitutability=self._drug_substitutability(drug),
                buffer_days=self._estimate_buffer_days(drug),
                strategic_priority=1.0 if drug.nlem_tier.value == "TIER_1" else 0.5,
            )

        # Add API nodes and collect unique provinces
        provinces = set()
        for api in self.data_loader.get_apis():
            self.graph.add_node(
                api.id,
                type="api",
                name=api.name,
                sector="pharma",
                china_share=api.china_share,
                primary_provinces=api.primary_provinces,
                monthly_import_value_usd_millions=api.monthly_import_value_usd_millions,
                substitutability=self._api_substitutability(api),
                buffer_days=max(5, int(30 * (1 - api.china_share))),
            )
            for p in api.primary_provinces:
                provinces.add(p)

        # Add Province nodes
        for p_name in provinces:
            self.graph.add_node(
                p_name,
                type="province",
                name=p_name,
                country="China",
                sector="multi",
            )

        # Add Edges with trade_volume_share
        for edge in self.data_loader.get_dependencies():
            self.graph.add_edge(
                edge.source,
                edge.target,
                weight=edge.weight,
                edge_type=edge.edge_type,
                trade_volume_share=edge.weight,  # normalized 0-1
            )

        # Add Rare Earth nodes (from rare_earths.json if available)
        self._add_rare_earth_nodes()

        logger.info(
            f"Graph built: {self.graph.number_of_nodes()} nodes, "
            f"{self.graph.number_of_edges()} edges"
        )

    def _add_rare_earth_nodes(self):
        """Adds rare earth mineral nodes and their province edges."""
        try:
            import json
            import os
            seed_dir = self.data_loader.seed_dir
            re_path = os.path.join(seed_dir, "rare_earths.json")
            if not os.path.exists(re_path):
                return
            with open(re_path, "r", encoding="utf-8") as f:
                rare_earths = json.load(f)

            for re in rare_earths:
                node_id = re["id"]
                self.graph.add_node(
                    node_id,
                    type="input",
                    name=re["name"],
                    sector="rare_earth",
                    china_share=re.get("china_share", 0.9),
                    primary_provinces=re.get("primary_provinces", []),
                    monthly_import_value_usd_millions=re.get("monthly_import_value_usd_millions", 0),
                    end_use_sectors=re.get("end_use_sectors", []),
                    substitutability=max(0.0, 1.0 - re.get("china_share", 0.9)),
                    buffer_days=re.get("india_stockpile_days", 15),
                    strategic_priority=1.0 if re.get("criticality", 0) > 0.8 else 0.7,
                    criticality=re.get("criticality", 0.5),
                )

                # Create edges to provinces
                for prov in re.get("primary_provinces", []):
                    if not self.graph.has_node(prov):
                        self.graph.add_node(
                            prov, type="province", name=prov,
                            country="China", sector="multi"
                        )
                    share = re.get("china_share", 0.9) / max(1, len(re.get("primary_provinces", [prov])))
                    self.graph.add_edge(
                        node_id, prov,
                        weight=share,
                        edge_type="input_from_province",
                        trade_volume_share=share,
                    )
        except Exception as e:
            logger.warning(f"Failed to load rare earth nodes: {e}")

    # ─── Louvain Community Detection ─────────────────────────────────────

    def _detect_communities(self):
        """
        Runs Louvain community detection on the undirected version of the graph.
        Identifies supply chain clusters — groups of nodes that are tightly
        interconnected and likely to co-propagate shocks.
        """
        try:
            import community as community_louvain

            self._undirected = self.graph.to_undirected()
            self._communities = community_louvain.best_partition(
                self._undirected, random_state=42
            )

            # Build community summaries
            community_members: Dict[int, List[str]] = {}
            for node_id, comm_id in self._communities.items():
                community_members.setdefault(comm_id, []).append(node_id)

            # Label each community by its dominant province or sector
            for comm_id, members in community_members.items():
                provinces = [m for m in members if self.graph.nodes[m].get("type") == "province"]
                if provinces:
                    self._community_labels[comm_id] = f"{', '.join(provinces[:2])} Supply Cluster"
                else:
                    sectors = [self.graph.nodes[m].get("sector", "unknown") for m in members[:3]]
                    self._community_labels[comm_id] = f"{'|'.join(set(sectors))} Cluster {comm_id}"

            logger.info(
                f"Louvain detected {len(community_members)} communities "
                f"across {len(self._communities)} nodes"
            )

        except ImportError:
            logger.warning("python-louvain not installed. Community detection disabled.")
            self._community_detection_enabled = False
            self._community_detection_method = "disabled_missing_dependency"
            self._communities = {n: 0 for n in self.graph.nodes()}
        except Exception as e:
            logger.warning(f"Community detection failed: {e}")
            self._community_detection_enabled = False
            self._community_detection_method = "disabled_runtime_error"
            self._communities = {n: 0 for n in self.graph.nodes()}

    def get_communities(self) -> Dict[str, Any]:
        """Returns community detection results for the frontend."""
        community_data: Dict[int, Dict[str, Any]] = {}

        for node_id, comm_id in self._communities.items():
            if comm_id not in community_data:
                community_data[comm_id] = {
                    "id": comm_id,
                    "label": self._community_labels.get(comm_id, f"Cluster {comm_id}"),
                    "members": [],
                    "provinces": [],
                    "drugs": [],
                    "inputs": [],
                    "size": 0,
                }
            entry = community_data[comm_id]
            node = self.graph.nodes[node_id]
            node_type = node.get("type", "unknown")
            entry["members"].append(node_id)
            entry["size"] += 1
            if node_type == "province":
                entry["provinces"].append(node_id)
            elif node_type == "drug":
                entry["drugs"].append(node_id)
            elif node_type in ("api", "input"):
                entry["inputs"].append(node_id)

        return {
            "community_detection_enabled": self._community_detection_enabled,
            "community_detection_method": self._community_detection_method,
            "total_communities": len(community_data),
            "communities": list(community_data.values()),
        }

    def get_node_community(self, node_id: str) -> Optional[int]:
        """Returns the community ID for a given node."""
        return self._communities.get(node_id)

    def get_community_amplifier(self, community_id: int) -> float:
        """
        Returns the community amplifier C_community.
        Larger, more interconnected communities amplify shocks more.
        """
        members = [n for n, c in self._communities.items() if c == community_id]
        if not members:
            return 1.0
        # Amplifier based on community density (more edges = more amplification)
        subgraph = self._undirected.subgraph(members) if self._undirected else self.graph.subgraph(members)
        density = nx.density(subgraph) if len(members) > 1 else 0
        return 1.0 + (density * 0.5)  # Range: 1.0 to 1.5

    # ─── Personalized PageRank ───────────────────────────────────────────

    def compute_pagerank(
        self,
        shocked_node: Optional[str] = None,
        alpha: float = 0.85,
    ) -> Dict[str, float]:
        """
        Computes Personalized PageRank from a shock origin node.

        IMPORTANT: Since graph edges point drug→api→province, we run PageRank
        on the REVERSED graph when propagating from a province. This makes
        the random walk flow province→api→drug, correctly modeling how supply
        shocks propagate upstream to find affected end products.

        Args:
            shocked_node: The node where the shock originates (e.g., "Hebei")
            alpha: Damping factor (default 0.85)
        """
        cache_key = shocked_node or "__global__"
        if cache_key in self._pagerank_cache:
            return self._pagerank_cache[cache_key]

        personalization = None
        if shocked_node and self.graph.has_node(shocked_node):
            personalization = {shocked_node: 1.0}

        # Use reversed graph so PageRank flows from province → api → drug
        # This correctly models shock propagation through the supply chain
        target_graph = self.graph.reverse(copy=False)

        try:
            scores = nx.pagerank(
                target_graph,
                alpha=alpha,
                weight="trade_volume_share",
                personalization=personalization,
                max_iter=200,
            )
            self._pagerank_cache[cache_key] = scores
            return scores
        except nx.PowerIterationFailedConvergence:
            logger.warning("PageRank failed to converge, using uniform distribution")
            n = self.graph.number_of_nodes()
            return {node: 1.0 / n for node in self.graph.nodes()}

    def compute_province_risk(self, province_id: str, sector: str = "pharma") -> float:
        """
        Calculates the aggregated risk score for a Chinese province.
        Risk is based on:
          - Concentration (HHI) of the products produced in this province.
          - Total monthly import value dependent on this province.
          - Strategic priority of the dependent products.
        """
        if not self.graph.has_node(province_id):
            return 0.0

        dependent_products = []
        # Find APIs/Inputs that are LOCATED_IN or supplied by this province
        # In our graph, edges are drug -> api -> province, so predecessors(province) are the APIs
        for product_id in self.graph.predecessors(province_id):
            node_data = self.graph.nodes.get(product_id, {})
            if sector == "both" or node_data.get("sector") == sector:
                dependent_products.append(product_id)

        if not dependent_products:
            return 0.0

        total_risk = 0.0
        for product_id in dependent_products:
            node_data = self.graph.nodes.get(product_id, {})
            import_value = node_data.get("monthly_import_value_usd_millions", 1.0)
            china_share = node_data.get("china_share", 0.5)
            # High import value + High China share = High Dependency Risk
            product_risk = (import_value * china_share)
            total_risk += product_risk

        # Normalize to 0-100 scale (heuristic normalization)
        normalized_risk = min(100.0, math.log1p(total_risk) * 15)
        return round(normalized_risk, 2)

    # ─── Combined Risk Formula ───────────────────────────────────────────

    def compute_combined_risk(
        self,
        shocked_node: str,
        sector: str = "pharma",
    ) -> Dict[str, Dict[str, Any]]:
        """
        Computes the full risk score using the combined formula:

        R_i = PR_i(shock) × (1 - S_i) × e^(-B_i/τ) × C_community

        Where:
          PR_i  = Personalized PageRank from shock origin (on reversed graph)
          S_i   = Substitutability (0-1)
          B_i   = Buffer stock in days
          τ     = Sector lead time (pharma=30d, rare_earth=60d)
          C     = Community amplifier

        Returns a dict of {node_id: {score, components}} for all affected nodes.
        """
        tau = SECTOR_TAU.get(sector, 30)
        pr_scores = self.compute_pagerank(shocked_node)

        # Find max PR score for normalization (excluding the shocked node itself)
        pr_values = [v for k, v in pr_scores.items() if k != shocked_node and v > 0]
        pr_max = max(pr_values) if pr_values else 1.0

        results = {}
        for node_id, pr_score in pr_scores.items():
            if node_id == shocked_node:
                continue  # Skip the shock origin itself

            node_data = self.graph.nodes.get(node_id, {})
            node_type = node_data.get("type", "")

            # Only score drugs, APIs, and inputs — not provinces
            if node_type not in ("drug", "api", "input"):
                continue

            substitutability = node_data.get("substitutability", 0.3)
            buffer_days = max(1, node_data.get("buffer_days", 15))
            community_id = self._communities.get(node_id, 0)
            community_amp = self.get_community_amplifier(community_id)

            # Normalize PR score relative to max (so top node → ~1.0)
            pr_normalized = pr_score / pr_max if pr_max > 0 else 0

            # The combined formula
            non_substitutability = 1.0 - substitutability
            buffer_decay = math.exp(-buffer_days / tau)
            risk_raw = pr_normalized * non_substitutability * buffer_decay * community_amp

            # Scale to 0-100
            risk_score_normalized = min(100.0, risk_raw * 100)

            if risk_score_normalized > 0.5:  # Include anything meaningful
                results[node_id] = {
                    "risk_score": round(risk_score_normalized, 2),
                    "components": {
                        "pagerank": round(pr_score, 6),
                        "pagerank_normalized": round(pr_normalized, 4),
                        "non_substitutability": round(non_substitutability, 3),
                        "buffer_decay": round(buffer_decay, 3),
                        "community_amplifier": round(community_amp, 3),
                        "community_id": community_id,
                        "community_label": self._community_labels.get(community_id, ""),
                        "buffer_days": buffer_days,
                        "substitutability": round(substitutability, 3),
                    },
                    "type": node_type,
                    "name": node_data.get("name", node_id),
                    "sector": node_data.get("sector", sector),
                }

        # Sort by risk score descending
        results = dict(
            sorted(results.items(), key=lambda x: x[1]["risk_score"], reverse=True)
        )

        return results

    # ─── Existing Methods (Preserved) ────────────────────────────────────

    def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
        """Returns node data in a flat dictionary format."""
        if not self.graph.has_node(node_id):
            return None

        node_data = self.graph.nodes[node_id]
        return {
            "id": node_id,
            "type": node_data.get("type"),
            "name": node_data.get("name"),
            "attributes": {k: v for k, v in node_data.items() if k not in ["type", "name"]},
        }

    def get_neighbors(
        self,
        node_id: str,
        direction: Literal["in", "out", "both"] = "out",
    ) -> List[str]:
        """Returns the IDs of neighboring nodes based on direction."""
        if not self.graph.has_node(node_id):
            return []

        if direction == "out":
            return list(self.graph.successors(node_id))
        elif direction == "in":
            return list(self.graph.predecessors(node_id))
        else:
            return list(nx.all_neighbors(self.graph, node_id))

    def get_dependency_chain(self, drug_id: str) -> List[List[str]]:
        """
        Finds all simple paths of length 1-3 starting from a drug.
        Typically captures drug -> api -> province paths.
        """
        if not self.graph.has_node(drug_id):
            return []

        all_paths = []
        # Find paths up to length 3 (Drug -> API -> Province)
        for target in self.graph.nodes():
            if target == drug_id:
                continue
            # Use short path limits to capture the direct supply chain
            paths = list(nx.all_simple_paths(self.graph, source=drug_id, target=target, cutoff=3))
            all_paths.extend(paths)

        return all_paths

    def compute_concentration_hhi(self, drug_id: str) -> float:
        """
        Calculates the Herfindahl-Hirschman Index (HHI) for a drug's geographic dependency.
        Ranges from 0 to 10000. High values indicate extreme concentration.
        """
        if not self.graph.has_node(drug_id):
            return 0.0

        # 1. Map drug to provinces with effective weights
        province_weights: Dict[str, float] = {}

        # Get APIs
        for api_id in self.graph.successors(drug_id):
            edge_da = self.graph[drug_id][api_id]
            w_da = edge_da.get("weight", 1.0)

            # Get Provinces for this API
            for prov_id in self.graph.successors(api_id):
                edge_ap = self.graph[api_id][prov_id]
                w_ap = edge_ap.get("weight", 1.0)

                # Cumulative dependency share on this province
                share = w_da * w_ap
                province_weights[prov_id] = province_weights.get(prov_id, 0.0) + share

        if not province_weights:
            return 0.0

        # 2. Calculate HHI
        total_weight = sum(province_weights.values())
        if total_weight == 0:
            return 0.0

        hhi = 0.0
        for weight in province_weights.values():
            percentage = (weight / total_weight) * 100
            hhi += percentage ** 2

        return hhi

    def get_drugs_dependent_on_province(self, province_id: str) -> List[str]:
        """
        Returns a list of drug IDs that are dependent on the given province
        by traversing the graph in reverse.
        """
        if not self.graph.has_node(province_id):
            return []

        dependent_drugs = set()
        # Traverse backwards: Province <- API <- Drug
        for api_id in self.graph.predecessors(province_id):
            for drug_id in self.graph.predecessors(api_id):
                node_data = self.graph.nodes[drug_id]
                if node_data.get("type") == "drug":
                    dependent_drugs.add(drug_id)

        # Also check rare earth inputs
        for input_id in self.graph.predecessors(province_id):
            node_data = self.graph.nodes.get(input_id, {})
            if node_data.get("type") == "input":
                dependent_drugs.add(input_id)

        return list(dependent_drugs)

    def nodes_by_type(self, node_type: str) -> List[Dict[str, Any]]:
        """Returns all nodes of a specific type (e.g., 'api', 'province')."""
        return [
            self.get_node(n)
            for n, d in self.graph.nodes(data=True)
            if d.get("type") == node_type
        ]

    def to_serializable_dict(self) -> Dict[str, Any]:
        """Converts the graph structure to a JSON-serializable dictionary."""
        nodes = []
        for n, d in self.graph.nodes(data=True):
            node_entry = {
                "id": n,
                "type": d.get("type"),
                "name": d.get("name"),
                "attributes": {k: v for k, v in d.items() if k not in ["type", "name"]},
            }
            # Add community info
            if n in self._communities:
                node_entry["community_id"] = self._communities[n]
            nodes.append(node_entry)

        edges = []
        for u, v, d in self.graph.edges(data=True):
            edges.append({
                "source": u,
                "target": v,
                "weight": d.get("weight"),
                "edge_type": d.get("edge_type"),
            })

        return {"nodes": nodes, "edges": edges}

    def apply_risk_scores(self, risk_map: Dict[str, float]) -> None:
        """
        Updates the 'current_risk' attribute on drug nodes based on a map of {drug_id: risk_score}.
        """
        for drug_id, risk in risk_map.items():
            if self.graph.has_node(drug_id):
                self.graph.nodes[drug_id]["current_risk"] = risk

    # ─── Private Helpers ─────────────────────────────────────────────────

    @staticmethod
    def _drug_substitutability(drug: Any) -> float:
        """Estimates substitutability score for a drug (0=no alternatives, 1=many)."""
        if drug.has_substitute:
            return 0.6
        if drug.nlem_tier.value == "TIER_1":
            return 0.1  # Critical, hard to replace
        return 0.3

    @staticmethod
    def _api_substitutability(api: Any) -> float:
        """Estimates substitutability for an API based on China share."""
        # Higher China share = fewer alternatives = lower substitutability
        return max(0.0, 1.0 - api.china_share)

    @staticmethod
    def _estimate_buffer_days(drug: Any) -> int:
        """Estimates buffer stock in days based on drug tier and population."""
        if drug.nlem_tier.value == "TIER_1":
            return 10  # High-demand, low buffer
        elif drug.nlem_tier.value == "TIER_2":
            return 20
        return 30

    def invalidate_cache(self):
        """Clears the PageRank cache (call after graph modifications)."""
        self._pagerank_cache.clear()