Trace a supply chain across linked databases

"""Dump the cross-DB exchange chain from Ginko through Agribalyse for one process.

Traces every Ginko activity reached (non-zero scaling) when computing impacts
for a root process, then for each reached activity lists its cross-DB
exchanges (no activityLinkId) with: exchange amount, unit,
consumer refAmount, target Agribalyse activity name + refAmount + refUnit.

Used to diagnose unit-conversion / scaling mismatches between VoLCA and
brightway for Ginko-organic-apple-style datasets.

Usage:
    cd volca-deploy/volca/examples
    uv run cross_db_chain/dump_cross_db_chain.py <db> <processId>

https://www.volca.run/docs/python/
"""

import argparse
import sys
from collections import deque

from volca import Client, VoLCAError

BASE_URL = "http://localhost:8080"

client = Client(base_url=BASE_URL)


def get_activity(db, pid):
    # Raw JSON via the escape hatch: the typed ActivityDetail hides the wire
    # fields this diagnostic needs (activityLinkId, per-exchange role/tag).
    return client.call("get_activity", db_name=db, process_id=pid)["activity"]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("db", help="root database name")
    parser.add_argument("processId", help="root process ID (actUUID_prodUUID)")
    parser.add_argument("--max-depth", type=int, default=4)
    parser.add_argument("--out", default="cross_db_chain.csv")
    args = parser.parse_args()

    # BFS over internal links, recording cross-DB exchanges at every level
    visited = set()
    queue = deque([(args.db, args.processId, 1.0, 0, "root")])  # (db, pid, weight, depth, path)
    rows = []
    dep_act_cache = {}  # cache for Agribalyse activities

    while queue:
        db, pid, weight, depth, path = queue.popleft()
        if (db, pid) in visited:
            continue
        visited.add((db, pid))

        try:
            act = get_activity(db, pid)
        except VoLCAError as e:
            print(f"  [skip] {db}/{pid[:16]}: {e}", file=sys.stderr)
            continue

        consumer_ref_amount = act["productAmount"]
        consumer_ref_unit = act["productUnit"]
        consumer_name = act["activityName"]

        if depth >= args.max_depth:
            continue

        for ex_entry in act["exchanges"]:
            ex = ex_entry["exchange"]
            # role Input = non-reference technosphere input (was isInput && !isReference)
            if ex.get("role") != "Input" or ex.get("tag") != "TechnosphereExchange":
                continue

            amount = ex["amount"]
            unit = ex_entry.get("unitName", "?")
            link_id = ex.get("activityLinkId")
            target_pid = ex_entry.get("targetProcessId") or ""

            # Normalize consumer's per-kg coefficient
            per_kg_consumer = amount / consumer_ref_amount if consumer_ref_amount else amount

            if link_id is None and "::" in target_pid:
                # Cross-DB link
                dep_db, dep_pid = target_pid.split("::", 1)

                # Look up supplier refAmount/refUnit
                key = (dep_db, dep_pid)
                if key not in dep_act_cache:
                    try:
                        dep_act = get_activity(dep_db, dep_pid)
                        dep_act_cache[key] = {
                            "name": dep_act["activityName"],
                            "refAmount": dep_act["productAmount"],
                            "refUnit": dep_act["productUnit"],
                        }
                    except VoLCAError:
                        dep_act_cache[key] = {"name": "<missing>", "refAmount": None, "refUnit": "?"}
                dep_info = dep_act_cache[key]

                rows.append({
                    "depth": depth,
                    "path": path,
                    "consumer_db": db,
                    "consumer_name": consumer_name,
                    "consumer_ref_amount": consumer_ref_amount,
                    "consumer_ref_unit": consumer_ref_unit,
                    "exchange_amount": amount,
                    "exchange_unit": unit,
                    "per_kg_consumer": per_kg_consumer,
                    "weight_kg_of_consumer_per_kg_root": weight,
                    "effective_demand_kg_supplier_per_kg_root": per_kg_consumer * weight,
                    "supplier_db": dep_db,
                    "supplier_name": dep_info["name"],
                    "supplier_ref_amount": dep_info["refAmount"],
                    "supplier_ref_unit": dep_info["refUnit"],
                })
            elif link_id is not None and "::" not in target_pid and target_pid:
                # Internal link - recurse with scaled weight
                sub_weight = weight * per_kg_consumer
                queue.append((db, target_pid, sub_weight, depth + 1, f"{path} -> {consumer_name}"))

    # Write CSV
    import csv
    fields = ["depth", "path", "consumer_db", "consumer_name", "consumer_ref_amount",
              "consumer_ref_unit", "exchange_amount", "exchange_unit", "per_kg_consumer",
              "weight_kg_of_consumer_per_kg_root", "effective_demand_kg_supplier_per_kg_root",
              "supplier_db", "supplier_name", "supplier_ref_amount", "supplier_ref_unit"]
    with open(args.out, "w", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for row in sorted(rows, key=lambda r: -r["effective_demand_kg_supplier_per_kg_root"]):
            w.writerow(row)
    print(f"Wrote {len(rows)} cross-DB link rows to {args.out}")

    # Top 20 by effective demand
    print()
    print(f"{'supplier':70s}  {'amount':>12s}  {'unit':>6s}  {'cons.ref':>10s} {'cons.unit':>10s}  {'eff.dem':>12s}")
    for r in sorted(rows, key=lambda r: -abs(r["effective_demand_kg_supplier_per_kg_root"]))[:20]:
        print(f"{r['supplier_name'][:70]:70s}  "
              f"{r['exchange_amount']:>12.4g}  "
              f"{r['exchange_unit']:>6s}  "
              f"{r['consumer_ref_amount']:>10.1f}"
              f"{r['consumer_ref_unit']:>6s}    "
              f"{r['effective_demand_kg_supplier_per_kg_root']:>12.4g}")


if __name__ == "__main__":
    main()