json_to_sql.py

"""
Convert pH-aware hydrogen JSON files to a SQLite lookup database.
"""

import argparse
import json
import logging
import os
import sqlite3
import tempfile
import zlib
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)


# BLOB ENCODING

_ZLIB_LEVEL = 9 # Compression level, to make the db as light as possible
# Default values for empty blobs
_EMPTY_H_GEOM_HEX = zlib.compress(b"[]", _ZLIB_LEVEL).hex()  # Empty JSON list: no hydrogen geometries.
_EMPTY_ATOM_ROLES_HEX = zlib.compress(b"{}", _ZLIB_LEVEL).hex()  # Empty JSON object/dict: no atom roles.


def encode_h_geom(h_list: List[dict]) -> bytes:
    """
    Takes a list of H (where each H is represented with a dict) and converts it into a zlib-compressed JSON blob.
    Each hydrogen is stored as a positional array:
        [h_name, bonded_atom, bond_length, bond_angle_deg, angle_ref_atom, dihedral_deg, dihedral_ref_atom]
    """
    rows = [] # Compact representation of the H geometries; one row per H atom.
    for h in h_list:
        bl = round(h["bond_length"], 4) if h.get("bond_length") is not None else None # Extract rounded bond length
        ba = round(h["bond_angle_deg"], 4) if h.get("bond_angle_deg") is not None else None # Extract rounded bond angle
        dd = round(h["dihedral_deg"], 4) if h.get("dihedral_deg") is not None else None # Extract rounded dihedral angle       
        rows.append([
            h["h_name"],
            h["bonded_atom"],
            bl,
            ba,
            h["angle_ref_atom"],
            dd,
            h.get("dihedral_ref_atom"),
        ]) # [h_name, bonded_atom, bond_length, bond_angle_deg, angle_ref_atom, dihedral_deg, dihedral_ref_atom]
    return zlib.compress(json.dumps(rows, separators=(",", ":")).encode(), _ZLIB_LEVEL) # Converts rows to compact JSON
                                                                                        # Then encodes the JSON string as bytes
                                                                                        # Then compresses those bytes with zlib to produce the final BLOB


def encode_atom_roles(donors: Set[str], acceptors: Set[str],
                      cations: Set[str], anions: Set[str]) -> bytes:
    """
    Takes donors+acceptors+cations+anions string sets, and converts that into a zlib-compressed JSON blob.
    The blob contains a JSON object {atom_name: bitmask} where each bitmask encodes the atom's roles:
        bit 0 (1) = H-bond donor
        bit 1 (2) = H-bond acceptor
        bit 2 (4) = cationic
        bit 3 (8) = anionic
    """
    all_atoms = donors | acceptors | cations | anions # Set of all atoms that appear in at least one of the individual sets
    roles: Dict[str, int] = {} # Roles dictionary, to fill with {atom_name: bitmask}
    for atom in sorted(all_atoms): 
        roles[atom] = ( # Bitwise OR
            (1 if atom in donors    else 0) 
            | (2 if atom in acceptors else 0)
            | (4 if atom in cations  else 0)
            | (8 if atom in anions   else 0)
        )
    return zlib.compress(json.dumps(roles, separators=(",", ":")).encode(), _ZLIB_LEVEL)


def _remove_sqlite_artifacts(path: Path, *, include_main: bool) -> List[str]:
    """
    Remove SQLite-related files for a database path and return the paths actually deleted.
    If include_main is True, the main database file is removed too.
    Otherwise, only possible SQLite sidecar files are removed.
    """
    candidate_paths = [
        Path(str(path) + suffix)
        for suffix in ("-journal", "-wal", "-shm")
    ] # Candidate files to delete

    if include_main:
        candidate_paths.insert(0, path)

    removed = []
    for p in candidate_paths:
        if p.exists():
            p.unlink() # Remove the file in the p path
            removed.append(str(p)) # List of removed files

    return removed


def _existing_sqlite_sidecars(path: Path) -> List[Path]:
    """Return SQLite sidecar files currently present for a database path."""
    return [
        sidecar
        for suffix in ("-journal", "-wal", "-shm")
        if (sidecar := Path(str(path) + suffix)).exists()
    ]


def _make_temp_output_path(final_output: Path) -> Path:
    """Create a unique staged database path beside the final output."""
    fd, name = tempfile.mkstemp(
        prefix=f".{final_output.name}.",
        suffix=".tmp",
        dir=final_output.parent,
    )
    os.close(fd)
    return Path(name)


def _infer_grid_step(sorted_ph_values: List[float]) -> Optional[float]:
    """Infer the pH grid step from sorted unique pH values."""
    if len(sorted_ph_values) < 2:
        return None

    deltas = []
    for i in range(1, len(sorted_ph_values)):
        delta = sorted_ph_values[i] - sorted_ph_values[i - 1]
        if delta > 0:
            deltas.append(delta)
    if not deltas:
        return None
    return min(deltas)


def _encode_state_ranges(
    point_state_pairs: List[Tuple[float, int]],
    compound_id: str,
    grid_step: Optional[float],
) -> List[Tuple[float, float, int]]:
    """
    Collapse sorted (pH, state_num) points into contiguous same-state pH ranges.
    Points merge only when they have the same state and are adjacent on the pH grid. 
    """
    if not point_state_pairs:
        return [] # Gives an empty list if there are no pH/state pairs

    ranges: List[Tuple[float, float, int]] = [] # List of pH ranges for each state
    ph_start, current_state = point_state_pairs[0] # Marks where first range starts, and first state
    ph_end = ph_start # Initializing first range end
    closed_states: Set[int] = set()
    eps = 1e-9

    for pH, state_num in point_state_pairs[1:]:
        contiguous = (
            grid_step is None
            or abs((pH - ph_end) - grid_step) <= max(eps, abs(grid_step) * 1e-6)
        )
        if state_num == current_state and contiguous: # If it's the same contiguous state, extend the current range
            ph_end = pH
            continue
        ranges.append((ph_start, ph_end, current_state)) # Close the previous range
        closed_states.add(current_state)
        if state_num in closed_states:
            raise ValueError(
                f"State {state_num} appears in more than one separated pH range "
                f"for compound {compound_id}"
            )
        ph_start = pH # Start new range for new state
        ph_end = pH
        current_state = state_num

    ranges.append((ph_start, ph_end, current_state)) # Close last range
    return ranges


def create_database(db_path: Path) -> sqlite3.Connection:
    """Create the single-table SQLite lookup schema and return an open connection."""
    conn = sqlite3.connect(db_path) # Open a connection to the SQLite database, creating the file if needed
    journal_mode = conn.execute("PRAGMA journal_mode=DELETE").fetchone()[0]
    if journal_mode.lower() != "delete":
        conn.close()
        raise RuntimeError(f"Could not enable DELETE journal mode: {journal_mode}")
    cursor = conn.cursor() # Creates the cursor to execute SQL commands

    # Single lookup table: routes (compound_id, pH range) directly to blobs.
    # Query shape:
    #   SELECT h_geom, atom_roles, smiles
    #   FROM state_lookup
    #   WHERE compound_id = ? AND ? BETWEEN ph_min AND ph_max
    #   LIMIT 1
    cursor.execute(f"""
        CREATE TABLE IF NOT EXISTS state_lookup (
            state_pk    INTEGER PRIMARY KEY,
            compound_id TEXT NOT NULL,
            ph_min      REAL NOT NULL,
            ph_max      REAL NOT NULL,
            state_num   INTEGER NOT NULL,
            smiles      TEXT,
            h_geom      BLOB NOT NULL DEFAULT X'{_EMPTY_H_GEOM_HEX}',
            atom_roles  BLOB NOT NULL DEFAULT X'{_EMPTY_ATOM_ROLES_HEX}',
            CHECK (ph_max >= ph_min),
            UNIQUE(compound_id, state_num),
            UNIQUE(compound_id, ph_min)
        )
    """)

    conn.commit()
    return conn


def process_json_file(json_path: Path, conn: sqlite3.Connection) -> None:
    """Process a single JSON file and insert state/range rows into state_lookup."""
    logger.info(f"Processing {json_path.name}")  # Keeping track of the JSONs being processed

    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)  # JSON loaded as python dictionary

    compound_id = data["compound_id"]  # Compound ID is mandatory
    states = data.get("states", [])
    ph_map = data.get("ph_map", {})

    if not states:
        logger.warning(f"No states found for {compound_id}")
        return
    if not ph_map:
        logger.warning(f"No pH map found for {compound_id}")
        return

    # --- Collecting everything in Python before touching the DB ---
    state_by_num: Dict[int, dict] = {}  # We want to go from a list of states to a dictionary keyed by state
    for s in states:
        state_num = int(s["state_id"]) 
        if state_num in state_by_num:
            raise ValueError(
                f"State {state_num} is duplicated in states list for compound {compound_id}"
            )
        state_by_num[state_num] = s

    # pH float → state_num 
    all_ph_values: Set[float] = set()
    point_by_ph: Dict[float, int] = {}  # The int is the state number

    for ph_key, sid in ph_map.items():
        pH = float(ph_key)
        all_ph_values.add(pH)

        if sid in (None, "", "none"):  # In case there is no state id
            continue

        state_num = int(sid)
        if state_num not in state_by_num:  # Check taht we have information of that state
            raise ValueError(
                f"State {sid} referenced by ph_map for compound {compound_id} "
                "does not exist in states list"
            )
        point_by_ph[pH] = state_num

    sorted_point_state_pairs = sorted(point_by_ph.items(), key=lambda x: x[0])  # Sort (pH, state_num) pairs by pH
    grid_step = _infer_grid_step(sorted(all_ph_values))
    ph_ranges = _encode_state_ranges(sorted_point_state_pairs, compound_id, grid_step)  # Results in [(min, max, state), ...]

    # --- Encoding blobs and inserting one row per state pH range. --- 
    cursor = conn.cursor()
    total_h = 0  # H counter
    rows = []  # To accumulate the rows that we then insert in the db (compound_id, ph_min, ph_max, state_num, smiles, h_blob, r_blob)

    for ph_min, ph_max, state_num in ph_ranges:
        s = state_by_num[state_num]
        smiles = s.get("smiles_canonical")
        h_list = s.get("all_h", [])
        total_h += len(h_list)  # Add the number of H records in this state to the total
        h_blob = encode_h_geom(h_list)
        r_blob = encode_atom_roles(
            set(s.get("hbond_donors", [])),
            set(s.get("hbond_acceptors", [])),
            set(s.get("cationic_atoms", [])),
            set(s.get("anionic_atoms", [])),
        )
        rows.append((compound_id, ph_min, ph_max, state_num, smiles, h_blob, r_blob)) # One full db row

    cursor.executemany(
        """INSERT INTO state_lookup
           (compound_id, ph_min, ph_max, state_num, smiles, h_geom, atom_roles)
           VALUES (?, ?, ?, ?, ?, ?, ?)""",
        rows,
    )

    logger.info(
        "  %s: %d state(s), %d pH range(s), %d H record(s)",
        compound_id,
        len(states),
        len(ph_ranges),
        total_h,
    )


def main():
    parser = argparse.ArgumentParser(
        description="Convert pH-aware hydrogen JSON files to a single-table SQL lookup database"
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        required=True,
        help="Directory containing JSON files"
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("ph_hydrogen_dict.db"),
        help="Output SQLite database file (default: ph_hydrogen_dict.db)"
    )
    
    args = parser.parse_args()
    args.input_dir = args.input_dir.expanduser().resolve()

    if not args.input_dir.exists():
        logger.error(f"Input directory not found: {args.input_dir}")
        return 1

    json_files = sorted(args.input_dir.glob("*.json"))
    if not json_files:
        logger.error(f"No JSON files found in {args.input_dir}")
        return 1

    final_output = args.output.expanduser().resolve()
    final_output.parent.mkdir(parents=True, exist_ok=True)  # Creates the output folder if it doesn't exist
    temp_output = _make_temp_output_path(final_output)

    conn: Optional[sqlite3.Connection] = None
    failed_files = 0
    failed_names: List[str] = []  # Names of the files that failed

    try:
        logger.info("Creating staged database: %s", temp_output)
        conn = create_database(temp_output)

        # Process all JSON files
        logger.info(f"Found {len(json_files)} JSON files to process")

        conn.execute("BEGIN")
        for json_path in json_files: 
            try:
                conn.execute("SAVEPOINT one_file")  # Creates a rollback point before processing this JSON
                process_json_file(json_path, conn)
                conn.execute("RELEASE SAVEPOINT one_file")  # Releases the rollback point if the JSON succeeds
            except Exception as e:
                failed_files += 1
                failed_names.append(json_path.name)
                conn.execute("ROLLBACK TO SAVEPOINT one_file")
                conn.execute("RELEASE SAVEPOINT one_file")
                logger.error(f"Error processing {json_path.name}: {e}")
                import traceback
                traceback.print_exc()

        if failed_files:
            conn.rollback()
            raise RuntimeError(
                "Database build aborted because one or more JSON files failed to process"
            )

        conn.commit()
        logger.info("Running VACUUM for final compaction")
        conn.execute("VACUUM")
    except Exception as exc:
        if conn is not None:
            conn.close()
        removed_temp = _remove_sqlite_artifacts(temp_output, include_main=True)
        if removed_temp:
            logger.info(
                "Removed staged database artifacts after failure: %s",
                ", ".join(removed_temp),
            )
        if failed_files:
            logger.error(
                "Database build failed: %d/%d file(s) failed",
                failed_files,
                len(json_files),
            )
            logger.error("Failed files: %s", ", ".join(failed_names))
        else:
            logger.error("Database build failed before completion: %s", exc)
        logger.error("Existing database left unchanged: %s", final_output)
        return 2
    else:
        conn.close()

    try:
        staged_sidecars = _existing_sqlite_sidecars(temp_output)
        if staged_sidecars:
            raise RuntimeError(
                "Staged database still has SQLite sidecars after closing: "
                + ", ".join(map(str, staged_sidecars))
            )

        output_sidecars = _existing_sqlite_sidecars(final_output)
        if output_sidecars:
            raise RuntimeError(
                "Output database appears active or requires SQLite recovery; "
                "refusing to replace it while sidecars exist: "
                + ", ".join(map(str, output_sidecars))
            )
        os.replace(temp_output, final_output)
    except Exception as exc:
        removed_temp = _remove_sqlite_artifacts(temp_output, include_main=True)
        if removed_temp:
            logger.info(
                "Removed staged database artifacts after publish failure: %s",
                ", ".join(removed_temp),
            )
        logger.error("Failed to publish staged database: %s", exc)
        logger.error("Existing database left unchanged: %s", final_output)
        return 2

    logger.info(f"Database created successfully: {final_output}")

    # Print summary
    conn = sqlite3.connect(final_output)
    cursor = conn.cursor()
    
    cursor.execute("SELECT COUNT(DISTINCT compound_id) FROM state_lookup")
    num_compounds = cursor.fetchone()[0]

    cursor.execute("SELECT COUNT(*) FROM state_lookup")
    num_state_ranges = cursor.fetchone()[0]

    cursor.execute("SELECT COALESCE(SUM(LENGTH(h_geom)), 0) FROM state_lookup")
    h_geom_bytes = cursor.fetchone()[0]

    cursor.execute("SELECT COALESCE(SUM(LENGTH(atom_roles)), 0) FROM state_lookup")
    atom_roles_bytes = cursor.fetchone()[0]

    logger.info("\nDatabase summary:")
    logger.info(f"  Compounds:                {num_compounds}")
    logger.info(f"  State/range rows:         {num_state_ranges}")
    logger.info(f"  h_geom blobs total:       {h_geom_bytes / 1024:.1f} KB")
    logger.info(f"  atom_roles blobs total:   {atom_roles_bytes / 1024:.1f} KB")
    
    conn.close()

    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())