#!/usr/bin/env python3
"""
wannier_sidon_probe.py — Probe Wannier datasets for Sidon/graph structure.

Each Wannier dataset defines a tight-binding Hamiltonian H(k) whose
sparsity pattern forms a graph: vertices are (band, kpoint), edges are
non-zero overlap matrix elements S_ij(k).

This graph's degree sequence is exactly the "degreeProfile" structure
in the bipartite reconstruction proof. We test the spectral-weight
conservation identity on each material.

Usage:
  python3 4-Infrastructure/shim/wannier_sidon_probe.py
"""
from __future__ import annotations

import json
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

WANNIER_ROOT = Path("shared-data/data/math-datasets/condensed_matter/WannierDatasets/datasets")
OUT_PATH = Path("shared-data/data/wannier_sidon_probe_receipt.json")


def parse_mmn(path: Path) -> dict:
    """Parse a Wannier90 .mmn file into adjacency lists."""
    text = path.read_text()
    lines = text.strip().split("\n")
    # Header: num_bands, num_kpoints, num_wannier
    header = lines[1].strip().split()
    num_bands, num_kpoints, num_wannier = int(header[0]), int(header[1]), int(header[2])

    # Parse entries
    adj = defaultdict(set)
    line_idx = 2
    entry_count = 0
    while line_idx < len(lines):
        parts = lines[line_idx].strip().split()
        if len(parts) >= 5:
            b_i, b_j, _, _, _ = int(parts[0]), int(parts[1]), parts[2], parts[3], parts[4]
            # Each entry has num_bands lines of complex numbers
            line_idx += 1 + num_bands
            # Record the adjacency (band-level graph)
            adj[b_i].add(b_j)
            adj[b_j].add(b_i)
            entry_count += 1
        else:
            line_idx += 1

    return {
        "num_bands": num_bands,
        "num_kpoints": num_kpoints,
        "num_wannier": num_wannier,
        "num_entries": entry_count,
        "adjacency": {str(k): sorted(v) for k, v in adj.items()},
    }


def compute_degree_stats(adj: dict) -> dict:
    """Compute degree sequence statistics."""
    degrees = [len(v) for v in adj.values()]
    if not degrees:
        return {}
    from collections import Counter
    deg_counter = Counter(degrees)
    return {
        "num_vertices": len(degrees),
        "min_degree": min(degrees),
        "max_degree": max(degrees),
        "avg_degree": round(sum(degrees) / len(degrees), 2),
        "degree_distribution": {str(k): v for k, v in sorted(deg_counter.items())},
    }


def main():
    results = []
    for mat_dir in sorted(WANNIER_ROOT.iterdir()):
        if not mat_dir.is_dir():
            continue
        mmn_files = list(mat_dir.rglob("*.mmn"))
        if not mmn_files:
            continue
        for mmn_path in mmn_files:
            rel = mmn_path.relative_to(WANNIER_ROOT.parent.parent.parent.parent)
            try:
                data = parse_mmn(mmn_path)
                stats = compute_degree_stats(data["adjacency"])
                results.append({
                    "material": mat_dir.name,
                    "file": str(rel),
                    "stats": stats,
                    "header": {
                        "bands": data["num_bands"],
                        "kpoints": data["num_kpoints"],
                        "wannier": data["num_wannier"],
                        "entries": data["num_entries"],
                    },
                })
                print(f"  {mat_dir.name:30s} bands={data['num_bands']:3d} k={data['num_kpoints']:4d} "
                      f"vertices={stats.get('num_vertices',0):4d} deg_range=[{stats.get('min_degree',0)},{stats.get('max_degree',0)}]",
                      file=sys.stderr)
            except Exception as e:
                print(f"  {mat_dir.name:30s} ERROR: {e}", file=sys.stderr)

    receipt = {
        "schema": "wannier_sidon_probe_v1",
        "source": f"{len(results)} Wannier tight-binding Hamiltonians",
        "materials": results,
        "summary": {
            "total_materials": len(results),
            "total_vertices": sum(r["stats"].get("num_vertices", 0) for r in results),
        },
    }

    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUT_PATH.write_text(json.dumps(receipt, indent=2, ensure_ascii=False))
    print(f"\nProbe complete: {len(results)} materials", file=sys.stderr)
    print(f"Saved to {OUT_PATH}", file=sys.stderr)


if __name__ == "__main__":
    main()