Update page_files/Upload_Data.py

page_files/Upload_Data.py

@@ -23,7 +23,6 @@ import requests
 import streamlit as st
 from PIL import Image
 
-
 from dotenv import load_dotenv
 load_dotenv()
 
@@ -31,7 +30,9 @@ _GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 if not _GEMINI_API_KEY:
     raise RuntimeError("GEMINI_API_KEY not set in environment")
 
-# ── imports from doctodb_rag (data extraction) ────────────────────────────────
+# ── Data extraction (Code 1) ─────────────────────────────────────────────────
+# NOTE: DOI is injected locally below (see _extract_doi_from_pdf). If you want
+# the DOI produced inside PDF_DataExtraction itself, port the same regex there.
 from categorized.Backend.PDF_DataExtraction import (
     call_gemini_from_bytes,
     convert_to_dataframe,
@@ -39,29 +40,59 @@ from categorized.Backend.PDF_DataExtraction import (
     verify_dataframe,
 )
 
-# ── imports from figure_extractor (image extraction) ─────────────────────────
-from categorized.Backend.Pdf_ImageExtraction import (
-    GEMINI_MODEL as GEMINI_MODEL,
-    get_plot_data_from_llm,
-    extract_plots,
-)
-
+# ── New integrated stack: image2 → mapper5 → category_push ────────────────────
+# These are imported as top-level modules (the same way mapper5 imports image2
+# and category_push). Ensure image2.py / mapper5.py / category_push.py are on
+# the import path (same directory or added to sys.path above). If you keep them
+# under categorized/Backend/, change these three lines to:
+#   from categorized.Backend import image2, mapper5, category_push
+import image2
+import mapper5
+import category_push
+
+# Manual-entry path (unchanged — writes via the existing data_loader).
 from data_loader import insert_material_rows
-from categorized.Backend.plot_property_mapper import (
-    batch_map_plots,
-    fetch_properties_for_material,
-    save_plot_image_mapping,
-    save_plot_image_to_db,
-)
-from db import fetch_all
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Helpers that were previously in upload_backend
+# DOI extraction — regex over raw PDF text (format-verified, precision-first)
+# ─────────────────────────────────────────────────────────────────────────────
+
+_DOI_CORE_RE = re.compile(r'10\.\d{4,9}/[^\s"<>\]\)]+', re.IGNORECASE)
+
+
+def _clean_doi(doi: str) -> str:
+    if not doi:
+        return ""
+    doi = doi.strip()
+    doi = re.sub(r'^(?:https?://)?(?:dx\.)?doi\.org/', '', doi, flags=re.IGNORECASE)
+    doi = re.sub(r'^doi[:\s]+', '', doi, flags=re.IGNORECASE)
+    return doi.rstrip(' .,;:)]}>"\'')
+
+
+def _extract_doi_from_pdf(pdf_bytes: bytes) -> str:
+    """Prefer a doi.org URL, then a 'doi:' label, then any bare 10.xxxx/ token."""
+    try:
+        with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
+            text = "\n".join((p.get_text("text") or "") for p in doc)
+    except Exception:
+        return ""
+    m = re.search(r'(?:https?://)?(?:dx\.)?doi\.org/(10\.\d{4,9}/[^\s"<>\]\)]+)', text, re.IGNORECASE)
+    if m:
+        return _clean_doi(m.group(1))
+    m = re.search(r'\bdoi[:\s]+\s*(10\.\d{4,9}/[^\s"<>\]\)]+)', text, re.IGNORECASE)
+    if m:
+        return _clean_doi(m.group(1))
+    m = _DOI_CORE_RE.search(text)
+    return _clean_doi(m.group(0)) if m else ""
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Metadata helper
 # ─────────────────────────────────────────────────────────────────────────────
 
 def _df_to_meta(df: pd.DataFrame) -> dict:
-    """Re-create the flat metadata dict that the UI previously got from Gemini."""
+    """Re-create the flat metadata dict the UI expects."""
     if df.empty:
         return {}
     row0 = df.iloc[0]
@@ -71,151 +102,35 @@ def _df_to_meta(df: pd.DataFrame) -> dict:
         "material_abbreviation": str(row0.get("material_abbreviation", "")),
         "trade_grade":           str(row0.get("trade_grade", "")),
         "manufacturer":          str(row0.get("manufacturer", "")),
+        "doi":                   str(row0.get("doi", "")),
         "mechanical_properties": props,
     }
 
 
-def create_zip(image_results: list, include_json: bool = True) -> bytes:
-    """
-    Pack extracted plot images (and optional JSON metadata) into a ZIP.
-    Each item in image_results has: caption, page, image_data (list of dicts
-    with 'array' (BGR ndarray) and 'filename').
-    """
-    buf = io.BytesIO()
-    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
-        meta = []
-        for item in image_results:
-            caption = item.get("caption", "")
-            page    = item.get("page", "?")
-            for img_dict in item.get("image_data", []):
-                bgr      = img_dict.get("array")
-                filename = img_dict.get("filename", "plot.png")
-                if bgr is not None:
-                    ok, enc = cv2.imencode(".png", bgr)
-                    if ok:
-                        zf.writestr(filename, enc.tobytes())
-            if include_json:
-                meta.append({
-                    "caption":     caption,
-                    "page":        page,
-                    "image_count": len(item.get("image_data", [])),
-                    "images":      [d.get("filename") for d in item.get("image_data", [])],
-                })
-        if include_json and meta:
-            zf.writestr("metadata.json", json.dumps(meta, indent=4))
-    return buf.getvalue()
-
-
-def save_matched_images(
-    df: pd.DataFrame,
-    image_results: list,
-    save_dir: str = "images",
-) -> list:
-    """
-    Heuristically match extracted plot captions to property names in df and
-    save matched images to disk. Returns list of match-info dicts.
-    """
-    os.makedirs(save_dir, exist_ok=True)
-    saved = []
-    props = df["property_name"].str.lower().tolist() if "property_name" in df.columns else []
-
-    for item in image_results:
-        caption  = (item.get("caption") or "").lower()
-        best_prop = None
-        best_score = 0
-        for prop in props:
-            # simple overlap score: shared words
-            cap_words  = set(re.findall(r"\w+", caption))
-            prop_words = set(re.findall(r"\w+", prop))
-            score = len(cap_words & prop_words)
-            if score > best_score:
-                best_score = score
-                best_prop  = prop
-
-        if best_prop and best_score > 0:
-            for idx, img_dict in enumerate(item.get("image_data", [])):
-                bgr = img_dict.get("array")
-                if bgr is None:
-                    continue
-                safe_prop = re.sub(r"[^\w\-]", "_", best_prop)
-                filename  = f"{safe_prop}_{idx}.png"
-                filepath  = os.path.join(save_dir, filename)
-                cv2.imwrite(filepath, bgr)
-                saved.append({
-                    "property": best_prop,
-                    "caption":  item.get("caption", ""),
-                    "path":     filepath,
-                })
-    return saved
-
-
-def save_single_image_with_property(
-    bgr: np.ndarray,
-    property_name: str,
-    save_dir: str = "images",
-) -> str:
-    """Save a single BGR image tagged with a property name. Returns filepath."""
-    os.makedirs(save_dir, exist_ok=True)
-    safe = re.sub(r"[^\w\-]", "_", property_name)
-    filepath = os.path.join(save_dir, f"{safe}.png")
-    cv2.imwrite(filepath, bgr)
-    return filepath
-
-
 # ─────────────────────────────────────────────────────────────────────────────
-# extract_images adapter
-# Bridges figure_extractor's extract_plots API to the image_results list shape
-# expected by the rest of the UI (list of {caption, page, image_data}).
+# extract_images adapter — now a thin wrapper over image2.extract_and_verify_plots
+# Returns image2's native shape:
+#   [{caption, page, image_data:[{array, bytes, filename, subplot_label,
+#                                 subplot_caption, source, verification}]}]
+# which is exactly what mapper5.map_plots_to_properties and the display expect.
 # ─────────────────────────────────────────────────────────────────────────────
 
-_GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "AIzaSyBzyMFKEqcjsWpR-OGAY42T250o1O39v3Y")
-
 def extract_images(pdf_path: str) -> list:
-    """
-    Use figure_extractor to detect and crop plot images from a PDF path.
-    Returns a list compatible with the image_results shape used throughout the UI:
-      [{ "caption": str, "page": int, "image_data": [{"array": bgr_ndarray, "filename": str}] }]
-    """
     try:
-        # gemini_model = init_gemini(_GEMINI_API_KEY)
-        plot_data    = get_plot_data_from_llm( GEMINI_MODEL, pdf_path)
-        raw_plots    = extract_plots(
-            pdf_path=pdf_path,
-            plot_data=plot_data,
-            pad=22,
-            score_thresh=0.35,
+        with open(pdf_path, "rb") as f:
+            pdf_bytes = f.read()
+        plot_results, coverage = image2.extract_and_verify_plots(
+            pdf_bytes, verify_engines=["gemini"]
         )
+        # Drop captions naming photos/micrographs/logos and crops the verifier
+        # marked "discard" (keeps "keep"/"recrop"). Pure selection, non-mutating.
+        plot_results = mapper5.select_real_plots(plot_results)
+        st.session_state["plot_coverage"] = coverage
     except Exception as e:
         log.error(f"extract_images failed: {e}")
+        st.session_state["plot_coverage"] = {}
         return []
-
-
-
-
-    # raw_plots items: {caption, page, path, plot_score, plot_type}
-    # Convert to image_results shape
-    image_results = []
-    for item in raw_plots:
-        bgr = cv2.imread(item["path"]) if item.get("path") else None
-        # clean up temp file written by extract_plots
-        if item.get("path") and os.path.exists(item["path"]):
-            try:
-                os.remove(item["path"])
-            except Exception:
-                pass
-
-        page    = item.get("page", 1)
-        caption = item.get("caption", f"Figure (page {page})")
-        safe    = re.sub(r"[^\w\-]", "_", caption)[:40]
-        filename = f"page{page}_{safe}.png"
-
-        image_results.append({
-            "caption":    caption,
-            "page":       page,
-            "image_data": [{"array": bgr, "filename": filename}] if bgr is not None else [],
-        })
-
-    return image_results
+    return plot_results
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -415,20 +330,30 @@ def render_top_bar():
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Helpers for tab2 mapping UI
+# Helpers for the mapping UI
 # ─────────────────────────────────────────────────────────────────────────────
 
 def _confidence_badge(conf: str) -> str:
     colors = {"high": "#16a34a", "medium": "#d97706", "low": "#dc2626"}
     c = colors.get((conf or "low").lower(), "#6b7280")
-    return (
-        f"<span class='conf-badge' style='background:{c}'>"
-        f"{conf.upper()}</span>"
-    )
+    return f"<span class='conf-badge' style='background:{c}'>{(conf or '').upper()}</span>"
+
+
+def _score_confidence(score) -> str:
+    """Map a numeric match_score to a coarse confidence label for the badge."""
+    try:
+        s = float(score)
+    except (TypeError, ValueError):
+        return "low"
+    if s >= 0.75:
+        return "high"
+    if s >= 0.45:
+        return "medium"
+    return "low"
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Manual input form
+# Manual input form (unchanged — writes via data_loader.insert_material_rows)
 # ─────────────────────────────────────────────────────────────────────────────
 
 def input_form():
@@ -576,101 +501,18 @@ def input_form():
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Tab 1: Material Data
-# Uses run_pipeline from doctodb_rag instead of call_gemini_from_bytes
+# Tab 1: Material Data  (extraction + DOI + RDS category selection)
 # ─────────────────────────────────────────────────────────────────────────────
 
-# def render_material_data_tab(pdf_path: str):
-#     st.subheader("Material Properties Data")
-
-#     if not st.session_state.pdf_data_extracted:
-#         with st.spinner("Extracting material data…"):
-#             with open(pdf_path, "rb") as f:
-#                 pdf_bytes = f.read()
-
-            
-#             df, df_gemini, df_gpt, _chunks, api_errors, meta = run_pipeline(pdf_bytes)
-
-#             if api_errors:
-#                 for err in api_errors:
-#                     st.warning(err)
-
-#             if not df.empty:
-#                 # Build the metadata dict that the rest of the UI expects
-#                 data = _df_to_meta(df)
-#                 st.session_state.pdf_extracted_df   = df
-#                 st.session_state.pdf_data_extracted = True
-#                 st.session_state.pdf_extracted_meta = data
-#             else:
-#                 st.warning("No data extracted from PDF.")
-
-#     df = st.session_state.pdf_extracted_df
-
-#     if df.empty:
-#         return
-
-#     meta = st.session_state.get("pdf_extracted_meta", {})
-#     st.success(f"Extracted {len(df)} properties")
-
-#     col1, col2 = st.columns(2)
-#     col1.metric("Material",     meta.get("material_name",         "N/A"))
-#     col2.metric("Abbreviation", meta.get("material_abbreviation", "N/A"))
-
-#     st.dataframe(df, use_container_width=True, height=400)
-#     st.subheader("Assign Material Category")
-
-#     extracted_material_class = st.selectbox(
-#         "Select category for this material",
-#         ["Polymer", "Fiber", "Composite"],
-#         index=None,
-#         placeholder="Required before adding to database",
-#         key="tab1_material_class",
-#     )
-
-#     if st.button("+ Add to Database"):
-#         if not extracted_material_class:
-#             st.error("Please select a material category before adding.")
-#             return
-
-#         df["material_class"] = extracted_material_class
-#         df["material_type"]  = extracted_material_class
-
-#         if st.session_state.image_results:
-#             with st.spinner("Saving matched plot images…"):
-#                 saved_images = save_matched_images(
-#                     df, st.session_state.image_results, save_dir="images"
-#                 )
-#                 if saved_images:
-#                     st.success(f"Saved {len(saved_images)} plot image(s)")
-#                     with st.expander("View saved images"):
-#                         for img_info in saved_images:
-#                             st.write(f"**{img_info['property']}** → {img_info['caption']}")
-#                             st.write(f"Saved to: `{img_info['path']}`")
-#                 else:
-#                     st.info("No plots matched the extracted properties automatically.")
-
-#         st.session_state.setdefault("user_uploaded_data", pd.DataFrame())
-#         st.session_state["user_uploaded_data"] = pd.concat(
-#             [st.session_state["user_uploaded_data"], df], ignore_index=True
-#         )
-#         st.success(f"Added to {extracted_material_class} database!")
-# ── Stage labels and estimated durations for the progress display ─────────────
 _STAGE_LABELS = {
     0.00: ("Checking cache",           2),
-    0.05: ("Extracting tables & text", 15),
-    0.20: ("Extraction complete",      0),
-    0.25: ("Indexing into ChromaDB",   8),
-    0.40: ("Ranking chunks",           5),
-    0.50: ("Ranking complete",         0),
-    0.55: ("Building batches",         2),
-    0.60: ("Running Gemini + GPT-4o",  30),
-    0.90: ("Merging results",          3),
-    0.95: ("Consensus filtering",      4),
+    0.30: ("Extracting via Gemini",   20),
+    0.70: ("Verifying against source", 8),
     1.00: ("Done",                     0),
 }
 
-def _nearest_stage_label(pct: float) -> tuple[str, int]:
-    """Return (label, est_seconds_remaining) for the closest stage."""
+
+def _nearest_stage_label(pct: float):
     best_key = min(_STAGE_LABELS, key=lambda k: abs(k - pct))
     return _STAGE_LABELS[best_key]
 
@@ -679,28 +521,23 @@ def render_material_data_tab(pdf_path: str):
     st.subheader("Material Properties Data")
 
     if not st.session_state.pdf_data_extracted:
-
-        bar      = st.progress(0.0)
-        status   = st.empty()   # stage label + ETA
-        timer    = st.empty()   # elapsed clock
-
+        bar    = st.progress(0.0)
+        status = st.empty()
+        timer  = st.empty()
         start_ts = time.time()
 
         def _cb(msg: str, pct: float):
-            elapsed  = time.time() - start_ts
+            elapsed = time.time() - start_ts
             label, est_remaining = _nearest_stage_label(pct)
             bar.progress(min(pct, 1.0))
             status.markdown(
                 f"**{label}** &nbsp;·&nbsp; <span style='color:#64748b'>{msg}</span>",
                 unsafe_allow_html=True,
             )
-            if est_remaining > 0:
-                timer.caption(
-                    f"⏱ Elapsed: {elapsed:.0f}s &nbsp;·&nbsp; "
-                    f"Est. remaining: ~{est_remaining}s"
-                )
-            else:
-                timer.caption(f"⏱ Elapsed: {elapsed:.0f}s")
+            timer.caption(
+                f"⏱ Elapsed: {elapsed:.0f}s"
+                + (f" · Est. remaining: ~{est_remaining}s" if est_remaining > 0 else "")
+            )
 
         with open(pdf_path, "rb") as f:
             pdf_bytes = f.read()
@@ -708,37 +545,31 @@ def render_material_data_tab(pdf_path: str):
         _cb("Extracting via Gemini…", 0.30)
         data = call_gemini_from_bytes(pdf_bytes)
         df = convert_to_dataframe(data)
-        api_errors = []
+
+        # DOI — regex over the PDF text, attached to every row. doi_url is the
+        # column category_push preserves into <table>_extras.
+        doi = _extract_doi_from_pdf(pdf_bytes)
+        if not df.empty:
+            df["doi"] = doi
+            df["doi_url"] = doi
 
         if not df.empty:
             _cb("Verifying against source PDF…", 0.70)
             sentences = _extract_sentences(pdf_bytes)
             df = verify_dataframe(df, sentences)
 
-        row0 = df.iloc[0] if not df.empty else {}
-        meta = {
-            "material_name": str(row0.get("material_name", "")) if not df.empty else "",
-            "material_abbreviation": str(row0.get("material_abbreviation", "")) if not df.empty else "",
-        }
         _cb("Done.", 1.0)
         elapsed_total = time.time() - start_ts
         bar.progress(1.0)
         status.empty()
         timer.empty()
 
-        if api_errors:
-            for err in api_errors:
-                st.warning(err)
-
         if not df.empty:
-            data = _df_to_meta(df)
             st.session_state.pdf_extracted_df   = df
             st.session_state.pdf_data_extracted = True
-            st.session_state.pdf_extracted_meta = data
-            st.success(
-                f"✅ Extracted {len(df)} properties in {elapsed_total:.0f}s"
-                + (f" · {meta.get('batches', '?')} batch(es)" if meta.get('batches') else "")
-            )
+            st.session_state.pdf_extracted_meta = _df_to_meta(df)
+            st.session_state.pdf_doi            = doi
+            st.success(f"Extracted {len(df)} properties in {elapsed_total:.0f}s")
         else:
             st.warning("No data extracted from PDF.")
             return
@@ -748,444 +579,258 @@ def render_material_data_tab(pdf_path: str):
         return
 
     meta = st.session_state.get("pdf_extracted_meta", {})
+    doi  = st.session_state.get("pdf_doi", "")
 
-    col1, col2 = st.columns(2)
-    col1.metric("Material",     meta.get("material_name",         "N/A"))
-    col2.metric("Abbreviation", meta.get("material_abbreviation", "N/A"))
+    c1, c2, c3 = st.columns(3)
+    c1.metric("Material",     meta.get("material_name",         "N/A"))
+    c2.metric("Abbreviation", meta.get("material_abbreviation", "N/A"))
+    c3.metric("DOI",          doi or "—")
 
     st.dataframe(df, use_container_width=True, height=400)
-    st.subheader("Assign Material Category")
 
-    extracted_material_class = st.selectbox(
-        "Select category for this material",
-        ["Polymer", "Fiber", "Composite"],
+    st.subheader("Assign Material Category")
+    st.selectbox(
+        "Category table (routes to the RDS table on push)",
+        category_push.CATEGORY_TABLES,   # Composites_materials / Fibers / Polymers
         index=None,
-        placeholder="Required before adding to database",
-        key="tab1_material_class",
+        placeholder="Required before pushing to the database",
+        key="push_category",
+        help="This is the RDS table the mapped rows are pushed into (from the "
+             "Extracted Plots tab).",
     )
+    if st.session_state.get("push_category"):
+        st.caption(
+            f"Category **{st.session_state['push_category']}** selected. "
+            "Go to the **Extracted Plots** tab to map figures and push to RDS."
+        )
 
-    if st.button("+ Add to Database"):
-        if not extracted_material_class:
-            st.error("Please select a material category before adding.")
-            return
 
-        df["material_class"] = extracted_material_class
-        df["material_type"]  = extracted_material_class
+# ─────────────────────────────────────────────────────────────────────────────
+# Tab 2: Extracted Plots → map (mapper5) → store (SQLite) → push (category_push)
+# ─────────────────────────────────────────────────────────────────────────────
 
-        if st.session_state.image_results:
-            with st.spinner("Saving matched plot images…"):
-                saved_images = save_matched_images(
-                    df, st.session_state.image_results, save_dir="images"
-                )
-                if saved_images:
-                    st.success(f"Saved {len(saved_images)} plot image(s)")
-                    with st.expander("View saved images"):
-                        for img_info in saved_images:
-                            st.write(f"**{img_info['property']}** → {img_info['caption']}")
-                            st.write(f"Saved to: `{img_info['path']}`")
-                else:
-                    st.info("No plots matched the extracted properties automatically.")
+_OUT_DIR = os.getenv("AIM_OUT_DIR", os.path.abspath("./aim_efrc_db"))
+_SOURCE_TABLE = "gemini_verified"   # provenance label for the stored rows
 
-        st.session_state.setdefault("user_uploaded_data", pd.DataFrame())
-        st.session_state["user_uploaded_data"] = pd.concat(
-            [st.session_state["user_uploaded_data"], df], ignore_index=True
-        )
-        st.success(f"Added to {extracted_material_class} database!")
 
-# ───────────────────────��─────────────────────────────────────────────────────
-# Tab 2: Extracted Plots + AI Property Mapping
-# Uses extract_images (adapter above) instead of upload_backend's version
-# ─────────────────────────────────────────────────────────────────────────────
+def _links_for_group(links_df: pd.DataFrame, gi: int) -> pd.DataFrame:
+    if links_df is None or links_df.empty or "group_idx" not in links_df.columns:
+        return pd.DataFrame()
+    return links_df[links_df["group_idx"] == gi]
+
 
 def render_plots_tab(pdf_path: str, paper_id: str):
     st.subheader("Extracted Plot Images & Property Mapping")
 
-   
+    # 1) Extract plots once (image2 detection + recovery + crop verification)
     if not st.session_state.pdf_processed:
         with st.spinner("Extracting plots from PDF…"):
-            st.session_state.image_results = extract_images(pdf_path)
+            st.session_state.plot_results  = extract_images(pdf_path)
             st.session_state.pdf_processed = True
             st.session_state.mapping_done  = False
+            st.session_state.links_df      = pd.DataFrame()
+            st.session_state.df_aug        = pd.DataFrame()
+            st.session_state.store         = None
 
-    image_results = st.session_state.image_results
-
-    if not image_results:
+    plot_results = st.session_state.plot_results
+    if not plot_results:
         st.warning("No plots found in this PDF.")
         return
 
-    has_data = not st.session_state.pdf_extracted_df.empty
+    df = st.session_state.pdf_extracted_df
+    has_data = not df.empty
+    n_imgs = sum(len(g.get("image_data", [])) for g in plot_results)
 
+    cov = st.session_state.get("plot_coverage", {}) or {}
     if has_data:
-        mat_abbr      = st.session_state.pdf_extracted_df.iloc[0]["material_abbreviation"]
-        property_list = st.session_state.pdf_extracted_df["property_name"].unique().tolist()
+        mat_abbr = df.iloc[0]["material_abbreviation"]
         st.info(
-            f"**{len(image_results)} plots** extracted  |  "
-            f"Material: **{mat_abbr}**  |  "
-            f"{len(property_list)} properties available for mapping"
+            f"**{len(plot_results)} figures / {n_imgs} crops** extracted  |  "
+            f"Material: **{mat_abbr}**  |  {df['property_name'].nunique()} properties  |  "
+            f"recovered: {cov.get('total_recovered', 0)}"
         )
     else:
-        st.warning(
-            "Extract material data in the **Material Data** tab first "
-            "to enable AI property mapping."
-        )
-
-    subtab_images, subtab_json = st.tabs(["🖼 Images & Mapping", "{ } JSON Preview"])
-
-    # ════════════════════════════════════════════════════════════════════════
-    with subtab_images:
-
-        col_img, col_json_dl, col_all = st.columns(3)
-        with col_img:
-            st.download_button(
-                "⬇ Images Only",
-                data=create_zip(image_results, include_json=False),
-                file_name=f"{paper_id}_images.zip",
-                mime="application/zip",
-                use_container_width=True,
-                key="dl_images",
-            )
-        with col_json_dl:
-            json_meta = [
-                {"caption": r["caption"], "page": r["page"],
-                 "image_count": len(r["image_data"])}
-                for r in image_results
-            ]
-            st.download_button(
-                "⬇ JSON",
-                data=json.dumps(json_meta, indent=4),
-                file_name=f"{paper_id}_metadata.json",
-                mime="application/json",
-                use_container_width=True,
-                key="dl_json",
-            )
-        with col_all:
-            st.download_button(
-                "⬇ Download All",
-                data=create_zip(image_results, include_json=True),
-                file_name=f"{paper_id}_complete.zip",
-                mime="application/zip",
-                use_container_width=True,
-                key="dl_all",
-            )
-
-        st.divider()
-
-        if has_data:
-            col_cls, col_btn = st.columns([0.45, 0.55])
-
-            with col_cls:
-                map_class = st.selectbox(
-                    "Material class for DB lookup",
-                    ["Polymer", "Fiber", "Composite"],
-                    key="mapping_material_class",
-                    help="Routes to the correct PostgreSQL table.",
-                )
-
-            with col_btn:
-                st.write("")
-                st.write("")
-                run_mapping = st.button(
-                    "🤖  Run AI Property Mapping",
-                    type="primary",
-                    disabled=st.session_state.get("mapping_done", False),
-                    use_container_width=True,
-                )
-
-            if run_mapping:
-                df           = st.session_state.pdf_extracted_df
-                mat_abbr     = df.iloc[0]["material_abbreviation"]
-                extracted_json = st.session_state.get("pdf_extracted_meta", {})
-
-                with st.spinner("Fetching properties from PostgreSQL…"):
-                    try:
-                        db_properties = fetch_properties_for_material(
-                            mat_abbr, map_class, fetch_all
-                        )
-                    except Exception as exc:
-                        st.error(f"DB error: {exc}")
-                        db_properties = []
-
-                if not db_properties:
-                    st.warning(
-                        f"No DB rows found for **{mat_abbr}** in the **{map_class}** table. "
-                        "Mapping will use all available properties from the extracted data."
-                    )
-
-                prog = st.progress(0, text="Starting…")
-
-                def _on_progress(i, total, caption):
-                    pct = int((i / max(total, 1)) * 100)
-                    prog.progress(pct, text=f"Mapping {i+1}/{total}: {caption[:55]}…")
-
-                with st.spinner("AI is analysing plots…"):
-                    mapped = batch_map_plots(
-                        image_results=image_results,
-                        extracted_json=extracted_json,
-                        db_properties=db_properties,
-                        progress_callback=_on_progress,
-                    )
-
-                prog.progress(100, text="Done ✓")
-                st.session_state.mapped_results = mapped
-                st.session_state.mapping_done   = True
-                st.success(f"✅  Mapped {len(mapped)} plots — review below.")
-                st.rerun()
-
-            if st.session_state.get("mapping_done"):
-                col_info, col_reset = st.columns([0.78, 0.22])
-                col_info.caption(
-                    "AI mapping complete. The dropdown for each plot is pre-filled "
-                    "with the suggestion — override freely, then hit **Save**."
-                )
-                if col_reset.button("↺ Re-run Mapping", use_container_width=True):
-                    st.session_state.mapping_done   = False
-                    st.session_state.mapped_results = []
-                    st.rerun()
+        st.warning("Extract material data in the **Material Data** tab first to enable mapping.")
+
+    # Downloads (image2's native zipper)
+    d1, d2 = st.columns(2)
+    d1.download_button(
+        "⬇ Images + metadata (ZIP)",
+        data=image2.create_plot_zip(plot_results, include_json=True),
+        file_name=f"{paper_id}_plots.zip", mime="application/zip",
+        use_container_width=True, key="dl_plots_zip",
+    )
+    d2.download_button(
+        "⬇ Images only (ZIP)",
+        data=image2.create_plot_zip(plot_results, include_json=False),
+        file_name=f"{paper_id}_images.zip", mime="application/zip",
+        use_container_width=True, key="dl_images_zip",
+    )
 
-        st.divider()
+    st.divider()
 
-        use_mapped = (
-            has_data
-            and st.session_state.get("mapping_done", False)
-            and bool(st.session_state.get("mapped_results"))
-        )
-        display_list = (
-            st.session_state.mapped_results if use_mapped else image_results
+    # 2) Map properties → figures (mapper5 four-signal cascade)
+    if has_data:
+        cA, cB = st.columns([0.6, 0.4])
+        run_map = cA.button(
+            "  Map properties → figures",
+            type="primary",
+            disabled=st.session_state.get("mapping_done", False),
+            use_container_width=True,
         )
+        if st.session_state.get("mapping_done"):
+            if cB.button("↺ Re-run mapping", use_container_width=True):
+                st.session_state.mapping_done = False
+                st.session_state.links_df = pd.DataFrame()
+                st.session_state.df_aug = pd.DataFrame()
+                st.session_state.store = None
+                st.rerun()
 
-        for idx in range(len(display_list)):
-            if idx >= len(display_list):
-                break
-
-            item     = display_list[idx]
-            caption  = item.get("caption", f"Figure {idx+1}")
-            page     = item.get("page", "?")
-            img_list = item.get("image_data", [])
-            mapping  = item.get("mapping_result") if use_mapped else None
-
-            with st.container(border=True):
-
-                col_cap, col_del = st.columns([0.87, 0.13])
-                col_cap.markdown(f"**Page {page}** — {caption}")
-                if col_del.button("🗑", key=f"del_grp_{idx}", help="Delete this figure"):
-                    display_list.pop(idx)
-                    if use_mapped:
-                        st.session_state.mapped_results = display_list
-                    else:
-                        st.session_state.image_results  = display_list
-                    st.rerun()
-
-                if mapping:
-                    prop_name  = mapping.get("property_name", "")
-                    section    = mapping.get("section", "")
-                    confidence = mapping.get("confidence", "low")
-                    reasoning  = mapping.get("reasoning", "")
-                    db_row     = mapping.get("db_row")
-                    candidates = mapping.get("all_candidates", [])
-
-                    if prop_name:
-                        badge = _confidence_badge(confidence)
-                        st.markdown(
-                            f"🔗 **AI Match:** `{section}` › **{prop_name}** &nbsp; {badge}",
+        if run_map:
+            with st.spinner("Fusing figure-citation → page → SciBERT → token signals…"):
+                links_df, df_aug = mapper5.map_plots_to_properties(df, plot_results)
+            st.session_state.links_df     = links_df
+            st.session_state.df_aug       = df_aug
+            st.session_state.mapping_done = True
+            n_mapped = int((df_aug.get("map_score", pd.Series(dtype=str)).astype(str) != "").sum()) \
+                if not df_aug.empty else 0
+            st.success(f" {n_mapped}/{len(df)} property rows linked to a figure "
+                       f"({len(links_df)} total link(s)).")
+            st.rerun()
+
+    links_df = st.session_state.get("links_df", pd.DataFrame())
+    mapping_done = st.session_state.get("mapping_done", False)
+
+    st.divider()
+
+    # 3) Figure-centric review — each figure with the property rows linked to it
+    for gi, group in enumerate(plot_results):
+        caption = group.get("caption", f"Figure {gi+1}")
+        page    = group.get("page", "?")
+        imgs    = group.get("image_data", [])
+
+        with st.container(border=True):
+            st.markdown(f"**Page {page}** — {caption}")
+
+            icols = st.columns(min(len(imgs), 4) or 1)
+            for pos, im in enumerate(imgs):
+                with icols[pos % len(icols)]:
+                    arr = im.get("array")
+                    if arr is not None:
+                        sub = im.get("subplot_label") or ""
+                        st.image(arr, channels="BGR", width=200,
+                                 caption=(sub or None))
+                    v = im.get("verification") or {}
+                    act = v.get("majority_action") if isinstance(v, dict) else None
+                    if act and act != "keep":
+                        st.caption(f"crop verdict: {act}")
+
+            if mapping_done:
+                grp = _links_for_group(links_df, gi)
+                if grp.empty:
+                    st.caption("No property linked to this figure.")
+                else:
+                    hdr = st.columns([3, 1.4, 1, 2, 1.2, 1, 0.7])
+                    for h, t in zip(hdr, ["property", "value", "unit", "material",
+                                          "subplot", "score", ""]):
+                        h.caption(t)
+                    for _, l in grp.iterrows():
+                        row = st.columns([3, 1.4, 1, 2, 1.2, 1, 0.7])
+                        row[0].write(str(l.get("property_name", "")))
+                        row[1].write(str(l.get("value", "")))
+                        row[2].write(str(l.get("unit", "")))
+                        row[3].write(str(l.get("material_name", "")))
+                        row[4].write(str(l.get("matched_subplot") or ""))
+                        row[5].markdown(
+                            _confidence_badge(_score_confidence(l.get("match_score"))),
                             unsafe_allow_html=True,
                         )
-                        if reasoning:
-                            st.caption(f"💬 {reasoning}")
-
-                        if db_row:
-                            with st.expander("📋 Matched DB row", expanded=False):
-                                c1, c2, c3 = st.columns(3)
-                                c1.metric("Value",     db_row.get("value",          "—"))
-                                c2.metric("Unit",      db_row.get("unit",           "—"))
-                                c3.metric("Condition", db_row.get("test_condition", "—"))
-                                if db_row.get("comments"):
-                                    st.caption(f"Comments: {db_row['comments']}")
-                                if db_row.get("english"):
-                                    st.caption(f"English units: {db_row['english']}")
-
-                        if candidates:
-                            with st.expander("🔄 All candidates", expanded=False):
-                                for c in candidates:
-                                    st.markdown(
-                                        f"{c.get('rank','?')}. `{c.get('section','?')}` › "
-                                        f"**{c.get('property_name','?')}** &nbsp; "
-                                        f"{_confidence_badge(c.get('confidence','low'))}",
-                                        unsafe_allow_html=True,
-                                    )
-                    else:
-                        st.warning("⚠️ AI could not match this plot to any DB property.")
-
-                for p_idx in range(len(img_list)):
-                    if p_idx >= len(item.get("image_data", [])):
-                        break
-
-                    img_data = item["image_data"][p_idx]
-                    bgr      = img_data.get("array")
-                    if bgr is None:
-                        continue
-
-                    img_key = f"{idx}_{p_idx}_{page}"
-                    st.image(bgr, channels="BGR", width=420)
-
-                    if has_data:
-                        df            = st.session_state.pdf_extracted_df
-                        mat_abbr      = df.iloc[0]["material_abbreviation"]
-                        property_list = df["property_name"].unique().tolist()
-                        options       = ["— Select property —"] + property_list
-
-                        ai_prop     = mapping.get("property_name", "") if mapping else ""
-                        ai_section  = mapping.get("section", "")        if mapping else ""
-                        default_idx = (
-                            property_list.index(ai_prop) + 1
-                            if ai_prop in property_list else 0
-                        )
-
-                        col_sel, col_sec, col_save, col_rem = st.columns(
-                            [0.40, 0.20, 0.20, 0.20]
-                        )
-
-                        with col_sel:
-                            selected = st.selectbox(
-                                "Property",
-                                options=options,
-                                index=default_idx,
-                                key=f"prop_sel_{img_key}",
-                                label_visibility="collapsed",
+                        key = f"rmlink_{gi}_{int(l.get('prop_row', 0))}_{int(l.get('match_rank', 0))}"
+                        if row[6].button("✕", key=key, help="Remove this link"):
+                            mask = ~(
+                                (links_df["group_idx"] == gi)
+                                & (links_df["prop_row"] == l["prop_row"])
+                                & (links_df["match_rank"] == l["match_rank"])
                             )
-
-                        with col_sec:
-                            section_options = [
-                                "Mechanical",
-                                "Thermal",
-                                "Processing",
-                                "Physical",
-                                "Descriptive",
-                                "Composition / Reinforcement",
-                                "Architecture / Structure",
-                            ]
-                            section_default = (
-                                section_options.index(ai_section)
-                                if ai_section in section_options
-                                else 0
-                            )
-                            section_val = st.selectbox(
-                                "Section",
-                                options=section_options,
-                                index=section_default,
-                                key=f"sec_{img_key}",
-                                label_visibility="collapsed",
+                            st.session_state.links_df = links_df[mask].reset_index(drop=True)
+                            # rebuild mapped_* columns from the pruned links
+                            st.session_state.df_aug = mapper5._apply_links(
+                                df, st.session_state.links_df
                             )
+                            st.rerun()
 
-                        with col_save:
-                            if st.button("💾 Save", key=f"save_{img_key}",
-                                         use_container_width=True):
-                                if selected and selected != "— Select property —":
-
-                                    filepath = save_plot_image_mapping(
-                                        mat_abbr, selected, section_val,
-                                        bgr, save_dir="images",
-                                    )
-
-                                    try:
-                                        from db import execute_query
-                                        saved_to_db = save_plot_image_to_db(
-                                            material_abbr=mat_abbr,
-                                            property_name=selected,
-                                            image_bgr=bgr,
-                                            material_class=st.session_state.get(
-                                                "mapping_material_class", "Polymer"
-                                            ),
-                                            execute_query_fn=execute_query,
-                                        )
-                                        if saved_to_db:
-                                            st.success(
-                                                f"✅ Saved to DB & disk → "
-                                                f"`{os.path.basename(filepath)}`"
-                                            )
-                                        else:
-                                            st.warning(
-                                                "⚠️ Saved to disk only — "
-                                                "no matching DB row found for this property."
-                                            )
-                                    except Exception as e:
-                                        st.error(f"DB save failed: {e}")
-                                        st.info(f"Saved locally → `{os.path.basename(filepath)}`")
-
-                                    st.session_state.saved_image_mapping[img_key] = {
-                                        "property": selected,
-                                        "section":  section_val,
-                                        "caption":  caption,
-                                        "filename": os.path.basename(filepath),
-                                        "path":     filepath,
-                                    }
-                                    st.rerun()
-                                else:
-                                    st.warning("Select a property first.")
-
-                        with col_rem:
-                            if st.button("✕", key=f"rem_{img_key}",
-                                         use_container_width=True, help="Remove image"):
-                                if img_key in st.session_state.saved_image_mapping:
-                                    del st.session_state.saved_image_mapping[img_key]
-                                item["image_data"].pop(p_idx)
-                                if not item["image_data"]:
-                                    display_list.pop(idx)
-                                    if use_mapped:
-                                        st.session_state.mapped_results = display_list
-                                    else:
-                                        st.session_state.image_results  = display_list
-                                st.rerun()
-
-                        if img_key in st.session_state.saved_image_mapping:
-                            saved_m = st.session_state.saved_image_mapping[img_key]
-                            st.info(
-                                f"✅ Saved as **{saved_m['property']}** → "
-                                f"`{saved_m['filename']}`"
-                            )
+    if not (has_data and mapping_done):
+        return
 
-                    else:
-                        col_msg, col_rem = st.columns([0.80, 0.20])
-                        col_msg.caption(
-                            "Go to **Material Data** tab to extract properties and enable mapping."
-                        )
-                        if col_rem.button("✕", key=f"rem_nd_{img_key}", help="Remove"):
-                            item["image_data"].pop(p_idx)
-                            if not item["image_data"]:
-                                st.session_state.image_results.pop(idx)
-                            st.rerun()
+    st.divider()
+
+    # 4) Store to SQLite (saves crops + attaches plot_image_path), then push to RDS
+    st.markdown("**Store the mapped result, then push to RDS**")
+    st.caption("Store writes the mapped rows + saved crops to a local SQLite DB and "
+               "attaches each crop's path. Push conforms those rows to the chosen "
+               "category table's 33-column schema and uploads them.")
+
+    if st.button(" Store to database (SQLite)", use_container_width=True):
+        with st.spinner("Saving crops and writing SQLite…"):
+            st.session_state.store = mapper5.store_properties_with_plots(
+                st.session_state.df_aug,
+                st.session_state.links_df,
+                plot_results,
+                out_dir=_OUT_DIR,
+                pdf_stem=paper_id,
+                source_table=_SOURCE_TABLE,
+            )
+        s = st.session_state.store
+        st.success(
+            f"Stored {s['n_properties']} row(s) → {s['db_filename']}; "
+            f"{s['n_properties_with_plot']} with a plot, {s['n_images_saved']} crop(s)."
+        )
+        st.rerun()
 
-                    st.divider()
-
-        saved_map = st.session_state.saved_image_mapping
-        if saved_map:
-            with st.expander(f"📁 Saved mappings ({len(saved_map)})", expanded=False):
-                for key, info in saved_map.items():
-                    st.markdown(
-                        f"**{info['property']}** &nbsp;›&nbsp; `{info['filename']}`  \n"
-                        f"<small style='color:#64748b'>Caption: {info['caption']}</small>",
-                        unsafe_allow_html=True,
-                    )
-
-    # ═════��══════════════════════════════════════════════════════════════════
-    with subtab_json:
-        st.subheader("Metadata Preview")
-        json_data = [
-            {
-                "caption":     r["caption"],
-                "page":        r["page"],
-                "image_count": len(r["image_data"]),
-                "images":      [img["filename"] for img in r["image_data"]],
-            }
-            for r in image_results
-        ]
-        st.download_button(
-            "⬇ Download JSON",
-            data=json.dumps(json_data, indent=4),
-            file_name="metadata.json",
-            mime="application/json",
-            key="dl_json_bottom",
+    store = st.session_state.get("store")
+
+    st.markdown("** Push to RDS database**")
+    if not mapper5.db_configured():
+        st.caption("Set DB_NAME / DB_USER / DB_PASSWORD in your `.env` "
+                   "(host is already configured) to enable the RDS push.")
+        return
+
+    category = st.session_state.get("push_category")
+    if not category:
+        category = st.selectbox(
+            "Category table (routes by material type)",
+            category_push.CATEGORY_TABLES,
+            key="push_category_tab2",
+            help="Rows conform to this table's 33-col schema; columns it can't hold "
+                 "go to <table>_extras; matched crops copy to rds_plots/<category>/<stem>/.",
         )
-        st.json(json_data)
+    embed_img = st.checkbox("Embed matched plot into the table's image column", value=True)
+
+    cT, cP = st.columns(2)
+    if cT.button(" Test connection", use_container_width=True):
+        try:
+            mapper5.db_healthcheck()
+            st.success("Connected to RDS.")
+        except Exception as e:
+            st.error(f"Connection failed: {e}")
+
+    push_ready = bool(store) and store.get("n_properties", 0) > 0
+    if cP.button(f" Push to {category}", type="primary",
+                 use_container_width=True, disabled=not push_ready):
+        try:
+            with st.spinner(f"Conforming + writing rows to '{category}'…"):
+                res = category_push.push_by_category(
+                    store, category, mapper5.get_db_engine(), embed_image=embed_img
+                )
+            st.success(
+                f"Pushed {res['pushed']} row(s) to '{res['table']}', "
+                f"{res['extras']} to '{res['extras_table']}', "
+                f"{res['plots_copied']} crop(s) → {res['plots_dir']}."
+            )
+        except Exception as e:
+            st.error(f"Push failed: {e}")
+    if not push_ready:
+        st.caption("Click ** Store to database (SQLite)** first — the push uploads those stored rows.")
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -1200,16 +845,19 @@ def main():
     st.caption("Provide technical data and research documentation for the central repository.")
 
     defaults = {
-        "image_results":       [],
-        "mapped_results":      [],
-        "pdf_processed":       False,
-        "mapping_done":        False,
-        "current_pdf_name":    None,
-        "form_submitted":      False,
-        "pdf_data_extracted":  False,
-        "pdf_extracted_df":    pd.DataFrame(),
-        "pdf_extracted_meta":  {},
-        "saved_image_mapping": {},
+        "plot_results":       [],
+        "plot_coverage":      {},
+        "links_df":           pd.DataFrame(),
+        "df_aug":             pd.DataFrame(),
+        "store":              None,
+        "pdf_processed":      False,
+        "mapping_done":       False,
+        "current_pdf_name":   None,
+        "form_submitted":     False,
+        "pdf_data_extracted": False,
+        "pdf_extracted_df":   pd.DataFrame(),
+        "pdf_extracted_meta": {},
+        "pdf_doi":            "",
     }
     for k, v in defaults.items():
         if k not in st.session_state:
@@ -1246,19 +894,13 @@ def main():
 
     if st.session_state.form_submitted:
         st.session_state.form_submitted = False
-        st.info(
-            "Form submitted. Previously extracted data has been saved. "
-            "Upload again to process a new PDF."
-        )
+        st.info("Form submitted. Upload again to process a new PDF.")
         st.tabs(["Material Data", "Extracted Plots"])
         return
 
-    tab1, tab2 = st.tabs(["📊 Material Data", "🖼 Extracted Plots"])
+    tab1, tab2 = st.tabs([" Material Data", "Extracted Plots"])
 
-    # Write to a stable temp file (avoids Windows WinError 267 on cleanup)
-    tmp_file = tempfile.NamedTemporaryFile(
-        suffix=".pdf", delete=False, prefix="matdb_"
-    )
+    tmp_file = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False, prefix="matdb_")
     try:
         tmp_file.write(uploaded_file.getbuffer())
         tmp_file.flush()
@@ -1267,10 +909,8 @@ def main():
 
         with tab1:
             render_material_data_tab(pdf_path)
-
         with tab2:
             render_plots_tab(pdf_path, paper_id)
-
     finally:
         try:
             os.unlink(tmp_file.name)
@@ -1278,5 +918,4 @@ def main():
             pass
 
 
-main()
-
+main()