import asyncio, re, textwrap, math, copy, os, uuid, json
from typing import Dict, List, Tuple
from datetime import datetime
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from PIL import Image
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from io import StringIO
from collections import defaultdict

from prompts import PROMPTS_CLASSIFICATION, PROMPTS_EXPERIMENTAL, PROMPTS_COMPUTATIONAL, PROMPTS_MIXED, PROMPTS_SYSTEM_CHART_NORMALIZATION

# -----helper functions and variables-----

_CATEGORY_COLORS = [
    '#E1BEE7',  # Light Purple
    '#F8BBD9',  # Light Pink
    '#FFCDD2',  # Light Red
    '#FFF9C4',  # Light Yellow
    '#B2DFDB',  # Light Teal
    '#D7CCC8',  # Light Brown
    '#CFD8DC',  # Light Blue Grey
    '#FFCCBC',  # Light Orange
    '#DCEDC8',  # Light Green
    '#B2EBF2'   # Light Cyan
]

def _clean_table_text(text) -> str:
    """
    Keep only the Markdown table in an LLM answer: drop code fences, blank
    lines, and any leading/trailing lines that are not table rows (prose such
    as "Here is the updated table:"). Table parsers below assume line 1 is the
    header, so a stray line on top would shift every row by one.
    """
    lines = [ln.rstrip() for ln in str(text or '').splitlines()]
    lines = [ln for ln in lines if ln.strip() and not ln.strip().startswith('```')]
    while lines and '|' not in lines[0]:
        lines.pop(0)
    while lines and '|' not in lines[-1]:
        lines.pop()
    return '\n'.join(lines)


def _is_separator_row(line: str) -> bool:
    """True for Markdown header-separator rows such as |---|:---:|---|."""
    return re.sub(r'[\|\-\:\s]', '', line) == '' and '-' in line and '|' in line


def _parse_microstructures(starter_table: str) -> str:
    """
    Extract the 'Structure' column (or 3rd real column if no headers) from a Markdown-like table.
    Handles both:
      1) No outer pipes, no explicit header separator
      2) With outer pipes and a '---|---|...' header separator
    Returns a '; '-joined string of non-empty cells from that column.
    """
    # 1) Keep only the table (no fences / prose), then normalize lines
    lines = [l.strip() for l in _clean_table_text(starter_table).splitlines() if l.strip()]
    if not lines:
        raise ValueError("Empty table text.")

    # 2) Remove obvious Markdown header separator lines (e.g., | --- | :---: | --- |)
    def is_md_separator(line: str) -> bool:
        # After removing pipes, colons, dashes and spaces, nothing should remain
        return re.sub(r'[\|\-\:\s]', '', line) == '' and ('-' in line or ':' in line) and '|' in line

    lines = [ln for ln in lines if not is_md_separator(ln)]

    # 3) Parse with pandas using '|' as delimiter
    text = '\n'.join(lines)
    df = pd.read_csv(StringIO(text), sep='|', engine='python', header=None, dtype=str, keep_default_na=False)

    # 4) Trim whitespace in all cells
    df = (df.map if hasattr(df, 'map') else df.applymap)(lambda x: x.strip() if isinstance(x, str) else x)

    # 5) Drop fully-empty boundary columns caused by leading/trailing pipes
    def col_empty(series):
        s = series.astype(str).str.strip()
        return s.eq('').all()

    while df.shape[1] > 0 and col_empty(df.iloc[:, 0]):
        df = df.iloc[:, 1:]
    while df.shape[1] > 0 and col_empty(df.iloc[:, -1]):
        df = df.iloc[:, :-1]

    if df.shape[1] < 1:
        raise ValueError("No columns detected after normalization.")

    # 6) Detect header row (case-insensitive). If first row contains 'structure', treat it as header.
    header_row = df.iloc[0].astype(str).tolist()
    header_lower = [h.strip().lower() for h in header_row]
    has_header = any(h == 'structure' for h in header_lower)

    if has_header:
        columns = [re.sub(r'^\*+(.*?)\*+$', r'\1', c or '').strip() for c in header_row]
        df = df.iloc[1:].reset_index(drop=True)
        df.columns = columns
        # Prefer the 'Structure' column by name; fall back to case-insensitive match
        if 'Structure' in df.columns:
            target = df['Structure']
        else:
            # case-insensitive match
            name_map = {c.lower(): c for c in df.columns}
            target = df[name_map.get('structure')]
    else:
        # No header: need at least 3 real columns, use the 3rd
        if df.shape[1] < 3:
            raise ValueError("Not enough columns to extract the 3rd column.")
        target = df.iloc[:, 2]

    # 7) Clean typical markdown bullets and noise
    def clean_cell(s: str) -> str:
        s = s.strip()
        s = re.sub(r'^\-\s+', '', s)          # leading "- "
        s = re.sub(r'\s*\(from text\)\s*$', '', s, flags=re.IGNORECASE)
        return s.strip('| ').strip()

    col = target.astype(str).map(clean_cell)

    # 8) Drop empty rows and join
    col = col[col.str.strip() != '']
    if col.empty:
        return ''

    # Deduplicate while keeping order (optional; drop if you want duplicates)
    seen = set()
    out = []
    for v in col:
        if v not in seen:
            seen.add(v)
            out.append(v)

    return '; '.join(out)

# ---- Paper category (Axiomatic design version) ----

PAPER_CATEGORIES = {'experimental': 'Experimental', 'computational': 'Computational', 'mixed': 'Mixed'}


def parse_paper_category(classification_result: str) -> str:
    """
    Parse 'Category: <X>' from the classification answer (Axiomatic design
    version: case-insensitive, text between 'Category:' and 'Reasons:').
    Returns one of 'Experimental', 'Computational', 'Mixed'; falls back to
    'Mixed' when the category cannot be parsed or is unknown.
    """
    text = str(classification_result or "")
    try:
        m = re.search(r"Category:\s*(.*?)\s*Reasons:", text, flags=re.DOTALL)
        raw = m.group(1) if m else text.split('Category: ')[1].splitlines()[0]
        key = raw.strip().strip('[]()*. ').lower()
        if key not in PAPER_CATEGORIES:
            print(f"Unknown paper category {raw.strip()!r}; falling back to Mixed.")
        return PAPER_CATEGORIES.get(key, 'Mixed')
    except Exception as e:
        print(f"There are something wrong with the classification result collection, specifically: {e}. Fallback to mixed case.")
        return 'Mixed'


# ---- Structure tagging from the step-1 categorization JSON (Axiomatic design version) ----

# Every tag_table_with_categories call records the constituent names it applied.
LAST_TAG_CATEGORIES: set = set()


def _extract_json_object(text: str) -> str:
    """
    Return the JSON object in `text` (tolerates ```json fences, prose, and
    other non-JSON blocks around it). Every '{' is tried as a start position
    with a strict JSON decoder; the first object that parses is returned,
    preferring one that carries a "categories" key.
    """
    text = str(text).strip()
    decoder = json.JSONDecoder()
    first_valid = None
    for m in re.finditer(r"\{", text):
        try:
            obj, _ = decoder.raw_decode(text, m.start())
        except ValueError:
            continue
        if isinstance(obj, dict):
            if "categories" in obj:
                return json.dumps(obj)
            if first_valid is None:
                first_valid = json.dumps(obj)
    return first_valid if first_valid is not None else text


def _normalize(text: str) -> str:
    text = text.lower()
    text = re.sub(r'\((?:[^()]|\([^()]*\))*\)', ' ', text)  # drop bracketed asides
    text = re.sub(r'[^a-z0-9\s]', ' ', text)                # keep alnum + space
    text = re.sub(r'\s+', ' ', text).strip()
    return text


def _match_score(row_text: str, target_text: str) -> float:
    """Hybrid score: containment bonus + token overlap."""
    r_norm = _normalize(row_text)
    t_norm = _normalize(target_text)
    r_tok, t_tok = set(r_norm.split()), set(t_norm.split())

    overlap = len(r_tok & t_tok) / max(1, len(t_tok))
    bonus = 0.0
    if t_norm and t_norm in r_norm:
        bonus = 1.0
    elif r_norm and r_norm in t_norm:
        bonus = 0.5
    return bonus + overlap


def tag_table_with_categories(starter_table: str, llm_response) -> str:
    """
    starter_table: markdown table string with columns:
      Processing | Mechanism (Processing & Structure) | Structure | Mechanism (Structure & Property) | Property
    llm_response: dict or JSON string with {"categories": {Category: [items...]}}
    returns: updated markdown table string with category tags appended in Structure column.

    Tagging policy:
      - Tag ALL occurrences of a microstructure in the table (not just one "best" row).
      - Prefer exact match on normalized text; fall back to fuzzy matching if needed.
    """
    # Load categories
    if isinstance(llm_response, str):
        llm = json.loads(_extract_json_object(llm_response))
    else:
        llm = llm_response
    cat_map: Dict[str, List[str]] = llm.get("categories", {})
    LAST_TAG_CATEGORIES.clear()
    LAST_TAG_CATEGORIES.update(k.strip() for k in cat_map if isinstance(k, str) and k.strip())

    # Split markdown into lines and basic parts (fences / prose removed first)
    lines = [ln.rstrip() for ln in _clean_table_text(starter_table).splitlines()]
    if len(lines) < 3:
        raise ValueError("Table appears to be missing header/separator/data rows.")

    header, separator = lines[0], lines[1]
    data_lines = lines[2:]

    # Parse data rows into list of columns (keep original spacing minimal)
    rows: List[List[str]] = []
    for ln in data_lines:
        cols = [c.strip() for c in ln.split('|')]
        # If the split produced leading/trailing empties, trim them
        if cols and cols[0] == '':
            cols = cols[1:]
        if cols and cols[-1] == '':
            cols = cols[:-1]
        if len(cols) < 5:
            # allow shorter, but we expect 5 columns; pad if needed
            cols += [''] * (5 - len(cols))
        rows.append(cols)

    # Index of the "Structure" column (0-based): Processing(0), Mech P->S(1), Structure(2), Mech S->P(3), Property(4)
    STRUCT_COL = 2

    # Build a cache of normalized structure cells for matching
    struct_norms = [_normalize(r[STRUCT_COL]) for r in rows]

    # Build list of (target_item, category) and a normalized lookup for exact matches
    targets: List[Tuple[str, str]] = []
    target_norm_to_categories: Dict[str, List[str]] = defaultdict(list)
    for category, items in cat_map.items():
        for it in items:
            if isinstance(it, str) and it.strip():
                t = it.strip()
                targets.append((t, category))
                t_norm = _normalize(t)
                if t_norm and category not in target_norm_to_categories[t_norm]:
                    target_norm_to_categories[t_norm].append(category)

    # Track categories per row (use sets during assignment; convert later)
    row_tag_sets: List[set] = [set() for _ in rows]

    # Pass 1: exact match on normalized text (tags all duplicates)
    for i, s_norm in enumerate(struct_norms):
        if not s_norm:
            continue
        for cat in target_norm_to_categories.get(s_norm, []):
            row_tag_sets[i].add(cat)

    # Pass 2: fuzzy fallback for any still-untagged rows
    # (use a slightly higher threshold to avoid over-tagging)
    FUZZY_THRESHOLD = 0.6
    for i, r in enumerate(rows):
        if row_tag_sets[i]:
            continue
        best_cat, best_score = None, 0.0
        for target, cat in targets:
            score = _match_score(r[STRUCT_COL], target)
            if score > best_score:
                best_score, best_cat = score, cat
        if best_cat is not None and best_score >= FUZZY_THRESHOLD:
            row_tag_sets[i].add(best_cat)

    # Convert sets to stable lists
    row_tags: List[List[str]] = [sorted(list(s)) for s in row_tag_sets]

    # Append tags to Structure cells (avoid double-tagging)
    for i, tags in enumerate(row_tags):
        if not tags:
            continue
        tag_str = " (" + ", ".join(tags) + ")"
        if not rows[i][STRUCT_COL].endswith(tag_str):
            rows[i][STRUCT_COL] = rows[i][STRUCT_COL] + tag_str

    # Rebuild the markdown table
    def join_row(cols: List[str]) -> str:
        return "| " + " | ".join(cols) + " |"

    new_lines = [header, separator] + [join_row(r) for r in rows]
    return "\n".join(new_lines)


_NA_VALUES = {'', 'n/a', 'na', 'n.a.', 'none', 'nil', 'null', 'not applicable', 'not available',
              'not specified', 'unknown', '-', '--', '–', '—'}


def _is_na(cell) -> bool:
    """True when a chart cell carries no information (N/A and its common variants).

    A trailing source tag such as "(From text)" or "(From knowledge base)" is
    ignored, so "N/A (From text)" counts as N/A.
    """
    text = re.sub(r'\s*\((?:from|source)[^)]*\)\s*$', '', str(cell or '').strip(), flags=re.IGNORECASE)
    text = text.strip().strip('*_` ').lower()
    return text in _NA_VALUES


def _wrap_label(text, wrap_limit=20):
    return textwrap.fill(text, width=wrap_limit)

def _parse_composition_and_key_components(response_text: str) -> tuple:
    """
    Parse composition and key added element from the LLM response.
    Supports variants like:
      - "Composition: ...; Key element addition: ..."   (current prompt)
      - "Composition: ...; Key Added Element: ..."
      - "Composition: ...; Key Components: ..."
      - "Composition: ...; Key Component: ..."
      - "Composition: ...; Key components addition: ..."
    Returns (composition, key_added_element). Falls back to "N/A" if missing.
    """
    if not response_text:
        return "N/A", "N/A"

    text = str(response_text).strip()
    composition = "N/A"
    key_components = "N/A"

    comp_label = re.search(r'Composition\s*:\s*', text, flags=re.IGNORECASE)
    key_label = re.search(r'Key\s+(?:Added\s+Element|Components?\s*(?:addition)?|[Ee]lement\s*(?:addition)?)\s*:\s*', text, flags=re.IGNORECASE)

    if comp_label:
        comp_start = comp_label.end()
        comp_end = key_label.start() if key_label and key_label.start() > comp_start else len(text)
        comp_value = text[comp_start:comp_end].strip().strip(';')
        composition = comp_value if comp_value else "N/A"

    if key_label:
        key_value = text[key_label.end():].strip().strip(';')
        key_components = key_value if key_value else "N/A"

    return composition, key_components

_NON_TEXT_BLOCK_TYPES = {'reasoning', 'function_call', 'tool_call', 'tool_use', 'tool_result',
                         'refusal', 'image', 'image_url', 'input_image', 'file', 'audio'}


def _llm_content_to_text(content) -> str:
    """
    Normalize heterogeneous LLM content payloads into a plain string.
    Handles str, list (content blocks), dict payloads, and fallback objects.
    """
    if content is None:
        return ""

    if isinstance(content, str):
        return content

    if isinstance(content, list):
        parts = []
        for item in content:
            if item is None:
                continue
            if isinstance(item, str):
                parts.append(item)
                continue
            if isinstance(item, dict):
                if isinstance(item.get('text'), str):
                    parts.append(item['text'])
                    continue
                text_block = item.get('text')
                if isinstance(text_block, dict):
                    value = text_block.get('value') or text_block.get('text')
                    if value is not None:
                        parts.append(str(value))
                        continue
                if isinstance(item.get('content'), str):
                    parts.append(item['content'])
                    continue
                # Non-text blocks (OpenAI Responses API reasoning items, tool calls,
                # refusals, ...) carry no answer text: skip them instead of dumping
                # their repr into the answer.
                if item.get('type') in _NON_TEXT_BLOCK_TYPES or 'encrypted_content' in item:
                    continue
                parts.append(str(item))
                continue

            text_attr = getattr(item, 'text', None)
            if isinstance(text_attr, str):
                parts.append(text_attr)
            else:
                parts.append(str(item))

        return '\n'.join([part for part in parts if part]).strip()

    if isinstance(content, dict):
        if isinstance(content.get('text'), str):
            return content['text']
        text_block = content.get('text')
        if isinstance(text_block, dict):
            value = text_block.get('value') or text_block.get('text')
            if value is not None:
                return str(value)
        if isinstance(content.get('content'), str):
            return content['content']
        return str(content)

    return str(content)

class _NormalizedResponse:
    """Proxy response object with normalized string `content`."""
    def __init__(self, response):
        self._response = response
        self.content = _llm_content_to_text(getattr(response, 'content', response))

    def __getattr__(self, name):
        return getattr(self._response, name)


class _NormalizedLLM:
    """LLM wrapper that normalizes `.ainvoke(...).content` to plain text."""
    def __init__(self, llm):
        self._llm = llm

    async def ainvoke(self, *args, **kwargs):
        response = await self._llm.ainvoke(*args, **kwargs)
        return _NormalizedResponse(response)

    def __getattr__(self, name):
        return getattr(self._llm, name)

def _transform_table_into_list(total_table):
    ps_list, sp_list = [], []
    lines = _clean_table_text(total_table).splitlines()
    for i in lines[1:]:  # Skip header; separator rows are skipped wherever they are
        if _is_separator_row(i):
            continue
        elements = [elem.strip() for elem in i.split('|') if elem]
        if len(elements) < 5:
            raise ValueError(f"Chart row has {len(elements)} cells, expected 5: {i[:100]!r}")
        ps_list.append((elements[0], elements[1], elements[2]))
        sp_list.append((elements[2], elements[3], elements[4]))
    return ps_list, sp_list

def _initialize_category_colors(G):
    """Extract categories and initialize colors for structure categories"""
    structure_categories = {attr['category'] for node, attr in G.nodes(data=True) if attr['group'] == 'structure'}
    category_colors = {}
    for i, category in enumerate(structure_categories):
        category_colors[category] = _CATEGORY_COLORS[i % len(_CATEGORY_COLORS)]
    return category_colors

def _get_text_width(text, fontsize=16):
    """Calculate the width of the wrapped text"""
    fig = plt.figure()
    renderer = fig.canvas.get_renderer()
    text_obj = plt.text(0, 0, text, fontsize=fontsize)
    bbox = text_obj.get_window_extent(renderer=renderer)
    plt.close(fig)
    return bbox.width / fig.dpi  # Convert from pixels to inches

def _create_main_diagram(G, pos, fig_width, fig_height, save_path, category_colors):
    """Create the main network diagram and save to file"""
    plt.figure(figsize=(fig_width, fig_height))  # Dynamic figsize
    ax = plt.gca()
    pos2 = copy.deepcopy(pos)

    dpi = ax.figure.get_dpi()

    # Draw nodes as rectangles
    key_list = list(pos.keys())
    for iii, nnn in enumerate(pos.items()):
        node, (x, y) = nnn
        group = G.nodes[node]['group']
        header = G.nodes[node].get('header', '')
        content = _wrap_label(G.nodes[node].get('content', node), wrap_limit=30)
        header_width = _get_text_width(header, fontsize=16)
        content_width = max([_get_text_width(line, fontsize=20) for line in content.split('\n')])
        width = content_width / 2.8
        height = 60. / dpi * (content.strip('\n').count('\n') + 1) + 0.6

        # Set color for nodes
        if group == 'structure':
            facecolor = category_colors.get(G.nodes[node]['category'], 'white')
        elif group == 'processing':
            facecolor = '#FFB74D'  # Orange
        elif group == 'property':
            facecolor = '#ADD8E6'  # Blue
        else:
            facecolor = 'white'

        rect = mpatches.FancyBboxPatch((x - width / 2, y - height / 2),
                                       width, height,
                                       boxstyle="square,pad=0.0",
                                       edgecolor='black',
                                       facecolor=facecolor)
        ax.add_patch(rect)
        plt.text(x, y + height / 2 + 36. / dpi / 2, header, fontsize=22, ha='center', va='center', fontweight='bold')
        plt.text(x, y, content, fontsize=22, ha='center', va='center')

        # Adjust node positions to avoid edges converging to the center
        if group == 'processing':
            pos[key_list[iii]] = (pos[key_list[iii]][0] + width / 2, pos[key_list[iii]][1])
            pos2[key_list[iii]] = (pos2[key_list[iii]][0] + width / 2, pos2[key_list[iii]][1])
        if group == 'property':
            pos[key_list[iii]] = (pos[key_list[iii]][0] - width / 2, pos[key_list[iii]][1])
            pos2[key_list[iii]] = (pos2[key_list[iii]][0] - width / 2, pos2[key_list[iii]][1])
        elif group == 'structure':
            pos[key_list[iii]] = (pos[key_list[iii]][0] + width / 2, pos[key_list[iii]][1])
            pos2[key_list[iii]] = (pos2[key_list[iii]][0] - width / 2, pos2[key_list[iii]][1])

    # Draw edges
    node_list = list(G.nodes())
    # Separate the edges based on their groups
    processing_edges = [(u, v) for u, v in G.edges() if G.nodes[u]['group'] == 'processing']
    property_edges = [(u, v) for u, v in G.edges() if G.nodes[v]['group'] == 'property']

    # Draw processing edges using pos2
    nx.draw_networkx_edges(G, pos2, edgelist=processing_edges, ax=ax)
    # Draw property edges using pos
    nx.draw_networkx_edges(G, pos, edgelist=property_edges, ax=ax)

    # Draw labels for edges
    for edge in G.edges():
        labels = G.edges[edge]['mechanisms']
        unique_labels = list(set(labels))
        combined_label = ', '.join(unique_labels)
        x, y = pos[edge[0]] if G.nodes[edge[0]]['group'] == 'processing' else pos[edge[0]]
        x2, y2 = pos[edge[1]] if G.nodes[edge[1]]['group'] == 'property' else pos2[edge[1]]
        xytext = ((x + x2) / 2, (y + y2) / 2)
        dx = (x2 - x) * 0.05
        dy = (y2 - y) * 0.05
        rot = math.degrees(math.atan2((y2 - y) * 0.23, (x2 - x)))
        plt.text(xytext[0] + dx, xytext[1] + dy, combined_label, size=18, ha='center', va='center', rotation=rot,
                bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', boxstyle='round,pad=0.3'))

    header_y_position = max(pos.values(), key=lambda x: x[1])[1] + 1
    plt.text(1, header_y_position + 1, 'Processing', fontsize=30, ha='center', va='center', fontweight='bold')
    plt.text(3, header_y_position + 1, 'Microstructure', fontsize=30, ha='center', va='center', fontweight='bold')
    plt.text(5, header_y_position + 1, 'Property', fontsize=30, ha='center', va='center', fontweight='bold')

    plt.tight_layout()
    plt.axis('off')
    plt.savefig(save_path, format='png', bbox_inches='tight')
    plt.close()
    return save_path

def _get_image_width(image_path):
    """Get the width of an image file"""
    with Image.open(image_path) as img:
        return img.width

def _add_na_nodes(G, mechanisms, pos):
    """Add N/A nodes to the graph if they exist in edges"""
    na_nodes = set()
    for edge in G.edges():
        if edge[0] == "N/A" or edge[1] == "N/A":
            na_nodes.update(edge)

    for na_node in na_nodes:
        if na_node == "N/A":
            continue
        if na_node not in G.nodes:
            G.add_node(na_node, group='na', header='N/A', content='N/A')
            pos[na_node] = (1, -1)

def _create_legend(mechanisms, category_colors, canvas_width_pixels, save_path):
    """Create a legend for the network diagram"""
    canvas_width_inches = canvas_width_pixels / 100
    fig_height = len(mechanisms) * 0.3 + len(category_colors) * 0.3
    fig, ax = plt.subplots(figsize=(canvas_width_inches, fig_height))
    legend_patches = [mpatches.Patch(color='none', label=f'{code}: {_wrap_label(mech, wrap_limit=150)}')
                      for mech, code in sorted(mechanisms.items(), key=lambda item: int(item[1])) if not _is_na(mech)]
    category_patches = [mpatches.Patch(color=color, label=f'{category}')
                        for category, color in category_colors.items()]
    ax.legend(handles=legend_patches + category_patches, loc='center', fontsize=22)
    plt.axis('off')
    plt.savefig(save_path, bbox_inches='tight')
    plt.close()
    return save_path

def _combine_images(image1_path, image2_path, output_path):
    """Combine two images vertically"""
    img1 = Image.open(image1_path)
    img2 = Image.open(image2_path)
    total_width = max(img1.width, img2.width)
    total_height = img1.height + img2.height
    new_img = Image.new('RGB', (total_width, total_height), 'white')
    new_img.paste(img2, ((total_width - img2.width) // 2, 0))
    new_img.paste(img1, ((total_width - img1.width) // 2, img2.height))
    new_img.save(output_path)
    return output_path

def _calculate_fig_dimensions(node_counts, max_label_length, min_height=10, height_per_node=2.0, width_per_char=0.2):
    """Calculate figure dimensions based on node counts and label length"""
    max_nodes = max(node_counts)
    fig_height = max(min_height, max_nodes * height_per_node)
    fig_width = max(20, max_label_length * width_per_char)
    return fig_width, fig_height

def _get_max_label_length(nodes):
    """Get the maximum label length from a list of nodes"""
    return max(len(label) for label in nodes)

def _split_microstructure_label(label):
    """Split microstructure label into name and category"""
    if '(' in label and ')' in label:
        parts = label.rsplit('(', 1)
        microstructure = parts[0].strip()
        category = parts[1].rstrip(')').strip()
        return microstructure, category
    return label, 'Other categories'

def _assign_unique_identifiers(ps_mechanisms, sp_mechanisms):
    """Assign unique identifiers to mechanisms"""
    unique_mechanisms = {}
    current_code = 1
    for mech in ps_mechanisms:
        if (not _is_na(mech)) and (mech not in unique_mechanisms):
            unique_mechanisms[mech] = str(current_code)
            current_code += 1
    for mech in sp_mechanisms:
        if (not _is_na(mech)) and (mech not in unique_mechanisms):
            unique_mechanisms[mech] = str(current_code)
            current_code += 1
    return unique_mechanisms

def _calculate_node_positions(nodes, x_position, vertical_gap, offset_y=0):
    """Calculate positions for nodes in the graph"""
    return {node: (x_position, offset_y + (len(nodes) - 1 - i) * vertical_gap) for i, node in enumerate(nodes)}

def create_network_diagram(system_chart, paper_title, temp_dir):
    """
    Main function to create network diagram from system chart
    
    Args:
        system_chart: The extracted system chart content
        paper_title: Title of the paper for filename generation
        temp_dir: Temporary directory to save the visualization files
    
    Returns:
        str: Path to the combined diagram file
    """
    import os
    import uuid
    import re
    
    ps_list, sp_list = _transform_table_into_list(system_chart.strip())

    # Rows with an N/A cell are not drawn, so they get no mechanism number either
    # (keeps the legend consistent with the edges actually shown).
    ps_list = [(p, m, s) for p, m, s in ps_list if not (_is_na(m) or _is_na(p) or _is_na(s))]
    sp_list = [(s, m, q) for s, m, q in sp_list if not (_is_na(m) or _is_na(s) or _is_na(q))]

    ps_mechanisms = [mech for _, mech, _ in ps_list]
    sp_mechanisms = [mech for _, mech, _ in sp_list]

    unique_mechanisms = _assign_unique_identifiers(ps_mechanisms, sp_mechanisms)

    G = nx.DiGraph()
    category_to_microstructures = defaultdict(set)

    for processing, mechanism, structure in ps_list:
        # Skip the edge when any of its cells is N/A: an "N/A" processing or
        # structure must not become a node.
        if not (_is_na(mechanism) or _is_na(processing) or _is_na(structure)):
            microstructure, category = _split_microstructure_label(structure)
            category_to_microstructures[category].add(microstructure)
            wrapped_processing = _wrap_label(processing)
            wrapped_microstructure = _wrap_label(microstructure)
            G.add_node(wrapped_processing, group='processing', category='Processing')
            G.add_node(wrapped_microstructure, group='structure', category=category)
            G.nodes[wrapped_microstructure]['header'] = ''
            G.nodes[wrapped_microstructure]['content'] = wrapped_microstructure
            if (wrapped_processing, wrapped_microstructure) in G.edges:
                G.edges[wrapped_processing, wrapped_microstructure]['mechanisms'].append(unique_mechanisms[mechanism])
            else:
                G.add_edge(wrapped_processing, wrapped_microstructure, mechanisms=[unique_mechanisms[mechanism]])

    for structure, mechanism, property in sp_list:
        if not (_is_na(mechanism) or _is_na(structure) or _is_na(property)):
            microstructure, category = _split_microstructure_label(structure)
            wrapped_property = _wrap_label(property)
            wrapped_microstructure = _wrap_label(microstructure)
            G.add_node(wrapped_microstructure, group='structure', category=category)
            G.nodes[wrapped_microstructure]['header'] = ''
            G.nodes[wrapped_microstructure]['content'] = wrapped_microstructure
            G.add_node(wrapped_property, group='property', category='Property')
            if (wrapped_microstructure, wrapped_property) in G.edges:
                G.edges[wrapped_microstructure, wrapped_property]['mechanisms'].append(unique_mechanisms[mechanism])
            else:
                G.add_edge(wrapped_microstructure, wrapped_property, mechanisms=[unique_mechanisms[mechanism]])

    # Define node positions before calling _add_na_nodes
    processing_nodes = [node for node, attr in G.nodes(data=True) if attr['group'] == 'processing']
    structure_nodes = [node for node, attr in G.nodes(data=True) if attr['group'] == 'structure']
    property_nodes = [node for node, attr in G.nodes(data=True) if attr['group'] == 'property']

    # Calculate the offsets to center the processing and property nodes
    processing_height = len(processing_nodes) * 3.0
    structure_height = len(structure_nodes) * 2.5
    property_height = len(property_nodes) * 3.0

    offset_y_processing = (structure_height - processing_height) / 2.5 if structure_height > processing_height else 0
    offset_y_property = (structure_height - property_height) / 2.5 if structure_height > property_height else 0

    processing_positions = _calculate_node_positions(processing_nodes, 1, 2.5, offset_y_processing)
    structure_positions = _calculate_node_positions(structure_nodes, 3, 2)
    property_positions = _calculate_node_positions(property_nodes, 5, 2.5, offset_y_property)

    pos = {**processing_positions, **structure_positions, **property_positions}

    max_label_length = _get_max_label_length(G.nodes)
    fig_width, fig_height = _calculate_fig_dimensions([len(processing_nodes), len(structure_nodes), len(property_nodes)], max_label_length)

    # Initialize category colors for structure nodes
    category_colors = _initialize_category_colors(G)

    # Add N/A nodes after defining pos
    _add_na_nodes(G, unique_mechanisms, pos)
    
    # Create safe filename from paper title
    safe_title = re.sub(r'[^\w\s-]', '', paper_title)[:50]  # Remove special chars, limit length
    safe_title = re.sub(r'[-\s]+', '_', safe_title)  # Replace spaces/hyphens with underscores
    unique_id = str(uuid.uuid4())[:8]
    
    # Create file paths
    diagram_path = os.path.join(temp_dir, f"diagram_{safe_title}_{unique_id}.png")
    legend_path = os.path.join(temp_dir, f"legend_{safe_title}_{unique_id}.png")
    combined_path = os.path.join(temp_dir, f"system_chart_{safe_title}_{unique_id}.png")
    
    # Create the diagrams
    _create_main_diagram(G, pos, fig_width, fig_height, diagram_path, category_colors)
    diagram_img_width = _get_image_width(diagram_path)
    _create_legend(unique_mechanisms, category_colors, diagram_img_width, legend_path)
    _combine_images(diagram_path, legend_path, combined_path)
    
    # Clean up intermediate files
    try:
        os.remove(diagram_path)
        os.remove(legend_path)
    except OSError:
        pass  # Don't fail if cleanup doesn't work
    
    return combined_path


def sanitize_filename(filename: str) -> str:
    """Sanitize filename by removing/replacing invalid characters"""
    import re
    # Remove or replace invalid characters
    filename = re.sub(r'[<>:"/\\|?*]', '_', filename)
    # Limit length
    return filename[:100] if len(filename) > 100 else filename

async def openai_api_function(input_list, llm):
    """
    Simple wrapper for OpenAI API calls using langchain
    
    Args:
        input_list: List of previous conversation messages
        llm: ChatOpenAI instance
    
    Returns:
        tuple: (updated input_list, latest response)
    """
    from langchain_core.messages import HumanMessage, AIMessage
    
    # Convert input_list to proper message format if needed
    messages = []
    for i, content in enumerate(input_list):
        if i % 2 == 0:  # Human messages (even indices)
            messages.append(HumanMessage(content=content))
        else:  # AI messages (odd indices)
            messages.append(AIMessage(content=content))
    
    # Get response from LLM
    response = await llm.ainvoke(messages)
    response_text = _llm_content_to_text(getattr(response, 'content', response))

    # Return updated list and response content
    return input_list, response_text

def _visualize_one_system_chart(system_chart_table_text: str, paper_text: str = "", citation: str = "", llm=None, temp_dir: str = "/tmp"):
    """
    Visualize system chart for a single paper, saving images in a citation-named subfolder.
    
    Args:
        system_chart_table_text: The system chart table as text
        paper_text: Full paper text for context
        citation: Citation/title for the paper
        llm: Language model instance
        temp_dir: Directory to save visualizations
    
    Returns:
        str: Path to the created visualization
    """
    import asyncio
    
    async def _async_visualize():
        print("Processing system chart...")
        
        # Parse the table to extract microstructures
        try:
            df = pd.read_csv(StringIO(system_chart_table_text), delimiter='|', skipinitialspace=True).iloc[1:, 1:-1].reset_index(drop=True)
            df.columns = df.columns.str.strip()
            microstructures = df['Microstructure'].str.strip().tolist()
            microstructures = '; '.join(microstructures)
        except Exception as e:
            print(f"Error parsing table: {e}")
            microstructures = "N/A"
        
        table_string = (
            "The table is as follows:\n"
            f" The table: {system_chart_table_text}"
        )
        
        input_list = []
        r_last = None
        
        # Use the normalization prompts as our chart-to-diagram prompts
        prompts_chart_to_diagram = PROMPTS_SYSTEM_CHART_NORMALIZATION
        
        for i in range(len(prompts_chart_to_diagram)):
            print(f"Processing step {i+1}/{len(prompts_chart_to_diagram)}...")
            
            if i == 0:
                prompt = prompts_chart_to_diagram[i].format(microstructures=microstructures)
                if paper_text:
                    prompt = (
                        f"{prompt}\n\n"
                        "-----\n"
                        "Full paper context (for reference only):\n"
                        f"{paper_text}"
                    )
            elif i == 1:
                prompt = prompts_chart_to_diagram[i] + table_string
            else:
                prompt = prompts_chart_to_diagram[i].format(r=r_last)
            
            input_list.append(prompt)
            
            if llm:
                input_list, r = await openai_api_function(input_list, llm)
                input_list.append(r)
                r_last = r
            else:
                print("Warning: No LLM provided, skipping API call")
                r_last = system_chart_table_text  # Use original table as fallback
        
        print("Creating diagrams...")
        
        # Save in visualized_system_charts/citation/
        base_folder = os.path.join(temp_dir, 'visualized_system_charts')
        os.makedirs(base_folder, exist_ok=True)
        
        safe_citation = sanitize_filename(citation) if citation else "unknown_citation"
        subfolder_path = os.path.join(base_folder, safe_citation)
        os.makedirs(subfolder_path, exist_ok=True)
        
        # Create the network diagram
        try:
            visualization_path = create_network_diagram(r_last or system_chart_table_text, citation, subfolder_path)
            print("Diagrams created!")
            return visualization_path
        except Exception as e:
            print(f"Error creating diagram: {e}")
            return None
    
    # Run the async function
    return asyncio.run(_async_visualize())

async def system_chart_extraction(paper_dict: dict, llm: ChatOpenAI, progress_callback=None, enable_visualization=False, temp_dir=None, extraction_id=None, run_normalization: bool = False) -> tuple:
    """
    Extract system charts from papers with detailed progress tracking.
    
    Args:
        paper_dict: Dictionary of title -> paper content
        llm: Language model for processing
        progress_callback: Function to call with progress updates
        enable_visualization: Whether to create visualizations
        temp_dir: Temporary directory to save visualization files
    """
    # Normalize LLM output at the boundary so downstream logic always receives strings.
    llm = _NormalizedLLM(llm)

    # Initialize output with header and format information
    output_system_charts = ''''''
    
    output_all_details = ''''''
    created_visualizations = []  # Track created visualization files
    
    total_papers = len(paper_dict)
    
    # Initialize progress
    if progress_callback:
        progress_callback({
            "status": "initializing",
            "progress": 0,
            "total_papers": total_papers,
            "current_paper": 0,
            "message": "Initializing..."
        })
    
    # First pass: classify all papers to calculate total questions
    paper_classifications = {}
    questions_per_type = {
        'Experimental': len(PROMPTS_EXPERIMENTAL),
        'Computational': len(PROMPTS_COMPUTATIONAL),
        'Mixed': len(PROMPTS_MIXED)
    }
    
    for idx, (title, paper) in enumerate(paper_dict.items()):
        if progress_callback:
            progress_callback({
                "status": "classifying paper type",
                "progress": int((idx / total_papers) * 20),  # Classification takes 20% of total
                "total_papers": total_papers,
                "current_paper": idx + 1,
                "message": f"Classifying paper {idx + 1}/{total_papers}: {title[:50]}..."
            })
        
        # Add small delay to allow frontend to catch up
        await asyncio.sleep(0.5)
        
        # Classify paper
        conversation_history = [HumanMessage(content=PROMPTS_CLASSIFICATION + '\nThe paper is:\n' + paper)]
        response = await llm.ainvoke(conversation_history)
        classification_result = response.content

        # Yield control back to event loop to allow other requests to be processed
        await asyncio.sleep(0)

        type_of_paper = parse_paper_category(classification_result)
            
        paper_classifications[title] = type_of_paper
    
    # Calculate total questions (classification + actual questions + normalization)
    total_questions = total_papers  # One classification per paper
    for paper_type in paper_classifications.values():
        total_questions += questions_per_type[paper_type]
    
    effective_normalization_steps = len(PROMPTS_SYSTEM_CHART_NORMALIZATION)
    if run_normalization:
        total_questions += effective_normalization_steps * total_papers  # Normalization steps
    
    completed_questions = total_papers  # We've completed all classifications
    
    # Update progress after classification phase is complete
    if progress_callback:
        progress_callback({
            "status": "classification_complete",
            "progress": 20,
            "total_papers": total_papers,
            "total_questions": total_questions,
            "completed_questions": completed_questions,
            "message": f"Classification complete. Starting to process {total_questions} total questions across {total_papers} papers..."
        })
    
    # Second pass: process each paper
    count_for_papers = 1
    
    for title, paper in paper_dict.items():
        paper_type = paper_classifications[title]
        
        # Called once per paper at the start of the paper processing, high level 
        if progress_callback:
            progress_callback({
                "status": "starting paper processing",
                "progress": int((completed_questions / total_questions) * 100),
                "total_papers": total_papers,
                "current_paper": count_for_papers,
                "total_questions": total_questions,
                "completed_questions": completed_questions,
                "current_paper_type": paper_type,
                "phase": "starting",
                "message": f"Starting paper {count_for_papers}/{total_papers} ({paper_type}): {title[:50]}..."
            })
        
        # Reset for each paper
        count_for_questions = 0
        count_for_visualization = 0
        conversation_history = []
        conversation_history_for_visualization = []
        five_answers_summary = {}
        
        # Add classification result to details
        output_all_details += ('\n')
        output_all_details += ('='*100 + '\n')
        output_all_details += (title)
        output_all_details += ('\n' + '='*100 + '\n')
        output_all_details += ('Check the type of the paper (Experimental, Computational or Mixed), this paper belongs to: ')
        output_all_details += (' '*5)
        output_all_details += (paper_type + '\n')
        output_all_details += ('\n')
        
        # Re-establish conversation history with classification
        conversation_history.append(HumanMessage(content=PROMPTS_CLASSIFICATION + '\n' + paper))
        classification_response = f"Category: {paper_type}\n"
        conversation_history.append(AIMessage(content=classification_response))
        
        def update_question_progress(question_num, if_visualization=False):
            nonlocal completed_questions
            completed_questions += 1
            if not if_visualization:
                if progress_callback:
                    progress_callback({
                        "status": "extracting system chart",
                        "progress": int((completed_questions / total_questions) * 100),
                        "total_papers": total_papers,
                        "current_paper": count_for_papers,
                        "total_questions": total_questions,
                        "completed_questions": completed_questions,
                        "current_paper_type": paper_type,
                        "current_question": question_num + 1,
                        "questions_in_paper": questions_per_type[paper_type],
                        "phase": "extraction",
                        "message": f"Paper {count_for_papers}/{total_papers} - Extracting system chart (Q{question_num + 1}/{questions_per_type[paper_type]})"
                    })
            else:
                if progress_callback:
                    # Provide specific status messages for different normalization steps
                    step_descriptions = {
                        0: "categorizing and tagging structures",
                        1: "standardizing properties",
                        2: "splitting mechanisms",
                        3: "separating multi-properties",
                        4: "filling missing values"
                    }
                    step_desc = step_descriptions.get(question_num, "normalizing system chart")
                    effective_normalization_steps = len(PROMPTS_SYSTEM_CHART_NORMALIZATION)
                    
                    progress_callback({
                        "status": step_desc,
                        "progress": int((completed_questions / total_questions) * 100),
                        "total_papers": total_papers,
                        "current_paper": count_for_papers,
                        "total_questions": total_questions,
                        "completed_questions": completed_questions,
                        "phase": "normalization",
                        "current_normalization_step": question_num + 1,
                        "total_normalization_steps": effective_normalization_steps,
                        "message": f"Paper {count_for_papers}/{total_papers} - Normalization: {step_desc} (Step {question_num + 1}/{effective_normalization_steps})"
                    })
        
        # Process questions based on paper type
        if paper_type == 'Experimental':
            prompts = PROMPTS_EXPERIMENTAL
            max_questions = len(PROMPTS_EXPERIMENTAL)
        elif paper_type == 'Computational':
            prompts = PROMPTS_COMPUTATIONAL
            max_questions = len(PROMPTS_COMPUTATIONAL)
        else:  # Mixed
            prompts = PROMPTS_MIXED
            max_questions = len(PROMPTS_MIXED)
        
        # Process all questions for this paper
        final_result = ""  # Initialize final_result
        while count_for_questions < max_questions:
            if paper_type == 'Experimental':
                r, sp, ps, output_all_details = await process_experimental_question(count_for_questions, prompts, conversation_history, 
                                                        llm, five_answers_summary, output_all_details, count_for_papers)
            elif paper_type == 'Computational':
                r, sp, ps, output_all_details = await process_computational_question(count_for_questions, prompts, conversation_history, 
                                                         llm, five_answers_summary, output_all_details, count_for_papers)
            else:  # Mixed
                r, sp, ps, output_all_details = await process_mixed_question(count_for_questions, prompts, conversation_history, 
                                                 llm, five_answers_summary, output_all_details, count_for_papers)
            
            # Store the result for the final system chart
            final_result = r
            
            # Update progress after each question
            update_question_progress(count_for_questions, if_visualization=False)
            
            count_for_questions += 1
        
        # Run normalization only when requested (visualization implies normalization).
        normalized_chart = None
        if final_result and run_normalization:
            # Notify that normalization phase is starting
            if progress_callback:
                progress_callback({
                    "status": "starting normalization",
                    "progress": int((completed_questions / total_questions) * 100),
                    "total_papers": total_papers,
                    "current_paper": count_for_papers,
                    "total_questions": total_questions,
                    "completed_questions": completed_questions,
                    "phase": "normalization_start",
                    "message": f"Paper {count_for_papers}/{total_papers} - Starting normalization..."
                })
            
            try:
                raw_chart = final_result
                normalized_chart = final_result
                print(f"\n🚀 Starting normalization process for paper: {title}")
                print(f"📊 Initial system chart length: {len(final_result)} characters")
                print("=" * 80)

                # Process all normalization steps (always execute, not just for visualization)
                effective_steps = len(PROMPTS_SYSTEM_CHART_NORMALIZATION)
                print(f"🔧 Total normalization steps to execute: {effective_steps}")
                
                while count_for_visualization < effective_steps:
                    print(f"\n⏳ Starting normalization step {count_for_visualization + 1}/{effective_steps}")
                    
                    normalized_chart = await normalize_system_charts(
                        count_for_visualization,
                        PROMPTS_SYSTEM_CHART_NORMALIZATION,
                        conversation_history_for_visualization,
                        normalized_chart,
                        paper,
                        llm
                    )
                    
                    print(f"📊 Chart length after step {count_for_visualization + 1}: {len(normalized_chart)} characters")
                    
                    # Update progress after each normalization step
                    update_question_progress(count_for_visualization, if_visualization=True)
                    count_for_visualization += 1

                print(f"\n✅ Normalization process completed! Final chart length: {len(normalized_chart)} characters")
                print("=" * 80)
                print("📋 FINAL NORMALIZED SYSTEM CHART:")
                print("-" * 60)
                print(normalized_chart[:1000] + "..." if len(normalized_chart) > 1000 else normalized_chart)
                print("-" * 60)
                print("=" * 80)

                # Keep the ORIGINAL final_result for system_charts.txt output
                # (normalized_chart is only used for visualization)
                
                # Optional: Create visualization if requested
                if enable_visualization and temp_dir:
                    visualization_path = await create_visualization(
                        normalized_chart,
                        title,
                        temp_dir,
                        created_visualizations
                    )
                    
                    if visualization_path:
                        print(f"✅ Visualization completed for paper: {title}")
                    else:
                        print(f"⚠️ Visualization creation failed for paper: {title}")
                    
            except Exception as e:
                print(f"❌ Normalization failed for paper '{title}': {str(e)}")
                # Continue processing other papers even if normalization fails
        
        # Add final result to system charts with new format
        delimiter = "=" * 100
        output_system_charts += f'\n{delimiter}\n'
        output_system_charts += f'PAPER: {title}\n'
        output_system_charts += f'SYSTEM_CHART: {final_result}\n'
        output_system_charts += f'{delimiter}\n\n'
        
        # Save conversation for this paper if extraction_id is provided
        if extraction_id:
            try:
                from conversation_storage import save_conversation
                conversation_key = save_conversation(
                    extraction_id=extraction_id,
                    paper_title=title,
                    conversation_history=conversation_history,
                    paper_content=paper,
                    paper_type=paper_type
                )
                print(f"💾 Saved conversation for paper: {title} (key: {conversation_key})")
            except Exception as e:
                print(f"⚠️ Could not save conversation for paper '{title}': {str(e)}")
        
        count_for_papers += 1
    
    # Final progress update
    if progress_callback:
        progress_callback({
            "status": "completed",
            "progress": 100,
            "total_papers": total_papers,
            "total_questions": total_questions,
            "completed_questions": completed_questions,
            "message": f"System chart extraction completed! Processed {total_papers} papers with {total_questions} total questions."
        })
    
    return output_system_charts, output_all_details, created_visualizations

async def process_experimental_question(count_for_questions, prompts, conversation_history, llm, 
                                five_answers_summary, output_all_details, count_for_papers):
    """Process a single question for experimental papers"""
    sp, ps = None, None
    r = ""
    qn = count_for_questions + 1
    
    if count_for_questions == 0:
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['property'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn} : Properties are: ' + r + '\n')
    elif count_for_questions == 1:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['microstructure'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: Microstructures are: '+ r + '\n')
    elif count_for_questions == 2:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        composition, key_elements = _parse_composition_and_key_components(r)
        five_answers_summary['composition&key_elements'] =  composition
        if 'N/A' not in key_elements:
            five_answers_summary['composition&key_elements'] += '; ' + key_elements
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: ' + r + '\n')
    elif count_for_questions == 3:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Enhanced properties of the alloy: {five_answers_summary['property']}; Microstructure of the alloy: {five_answers_summary['microstructure']}. "
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        sp = r
        five_answers_summary['sp_table'] = r
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 4:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['processing'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}' + ' Processing methods are: ' + r + '\n')
    elif count_for_questions == 5:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Processing methods of the alloy: {five_answers_summary['processing']}; Microstructure of the alloy: {five_answers_summary['microstructure']}. "
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        output_all_details += (f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 6:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Composition & Key components (if any) of the alloy: {five_answers_summary['composition&key_elements']}; Microstructure of the alloy: {five_answers_summary['microstructure']}. "
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        ps = r
        five_answers_summary['ps_table'] = r
        output_all_details += (f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 7:
        output_all_details += ('~'*100 + '\n')
        ps = five_answers_summary.get('ps_table')
        sp = five_answers_summary.get('sp_table')
        prerequisites = f"Prerequisites: First table for relationsip between processing and microstructure: {ps}; Second table for relationsip between microstructure and properties: {sp}. "
        conversation_history.append(HumanMessage(content=prerequisites + '\n' + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        output_all_details += ('\n' + f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    
    return r, sp, ps, output_all_details

async def process_computational_question(count_for_questions, prompts, conversation_history, llm, 
                                 five_answers_summary, output_all_details, count_for_papers=1):
    """Process a single question for computational papers"""
    sp, ps = None, None
    r = ""
    qn = count_for_questions + 1
    
    if count_for_questions == 0:
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['property'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn} : Properties are: ' + r + '\n')
    elif count_for_questions == 1:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['microstructure'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: Microstructures are: '+ r + '\n')
    elif count_for_questions == 2:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        composition, key_elements = _parse_composition_and_key_components(r)
        five_answers_summary['composition&key_elements'] =  composition
        if 'N/A' not in key_elements:
            five_answers_summary['composition&key_elements'] += '; ' + key_elements
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: ' + r + '\n')
    elif count_for_questions == 3:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Enhanced properties of the alloy: {five_answers_summary['property']}; Microstructure of the alloy: {five_answers_summary['microstructure']}."
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        sp = r
        five_answers_summary['sp_table'] = r
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 4:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites:Composition & Key components addition of the alloy: {five_answers_summary['composition&key_elements']};Microstructure of the alloy: {five_answers_summary['microstructure']}."
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        ps = r
        five_answers_summary['ps_table'] = r
        output_all_details += (f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 5:
        output_all_details += ('~'*100 + '\n')
        ps = five_answers_summary.get('ps_table')
        sp = five_answers_summary.get('sp_table')
        prerequisites = f"Prerequisites: First table for relationsip between processing and microstructure: {ps}; Second table for relationsip between microstructure and properties: {sp}. "
        conversation_history.append(HumanMessage(content=prerequisites + '\n' + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        output_all_details += ('\n' + f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    
    return r, sp, ps, output_all_details

async def process_mixed_question(count_for_questions, prompts, conversation_history, llm, 
                         five_answers_summary, output_all_details, count_for_papers):
    """Process a single question for mixed papers"""
    sp, ps = None, None
    r = ""
    qn = count_for_questions + 1
    
    if count_for_questions == 0:
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['property'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn} : Properties are: ' + r + '\n')
    elif count_for_questions == 1:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['microstructure'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: Microstructures are: '+ r + '\n')
    elif count_for_questions == 2:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        composition, key_elements = _parse_composition_and_key_components(r)
        five_answers_summary['composition&key_elements'] =  composition
        if 'N/A' not in key_elements:
            five_answers_summary['composition&key_elements'] += '; ' + key_elements
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: ' + r + '\n')
    elif count_for_questions == 3:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Enhanced properties of the alloy: {five_answers_summary['property']}; Microstructure of the alloy: {five_answers_summary['microstructure']}."
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        sp = r
        five_answers_summary['sp_table'] = r
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 4:
        output_all_details += ('~'*100 + '\n')
        conversation_history.append(HumanMessage(content=prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        five_answers_summary['processing'] = r.split(';')
        output_all_details += (f'Q{qn}: ' + prompts[count_for_questions] + '\n' + f'A{qn}' + ' Processing methods are: ' + r + '\n')
    elif count_for_questions == 5:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Processing methods of the alloy: {five_answers_summary['processing']}; Microstructure of the alloy: {five_answers_summary['microstructure']}. "
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        output_all_details += (f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 6:
        output_all_details += ('~'*100 + '\n')
        prerequisites = f"Prerequisites: Composition & Key components (if any) of the alloy: {five_answers_summary['composition&key_elements']}; Microstructure of the alloy: {five_answers_summary['microstructure']}. "
        conversation_history.append(HumanMessage(content=prerequisites + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        ps = r
        five_answers_summary['ps_table'] = r
        output_all_details += (f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    elif count_for_questions == 7:
        output_all_details += ('~'*100 + '\n')
        ps = five_answers_summary.get('ps_table')
        sp = five_answers_summary.get('sp_table')
        prerequisites = f"Prerequisites: First table for relationsip between processing and microstructure: {ps}; Second table for relationsip between microstructure and properties: {sp}. "
        conversation_history.append(HumanMessage(content=prerequisites + '\n' + prompts[count_for_questions]))
        response = await llm.ainvoke(conversation_history)
        conversation_history.append(AIMessage(content=response.content))
        await asyncio.sleep(0)
        r = response.content
        output_all_details += ('\n' + f'Q{qn}: ' + prerequisites + prompts[count_for_questions] + '\n' + f'A{qn}: '+ r + '\n')
    
    return r, sp, ps, output_all_details

async def normalize_system_charts(count_for_visualization, prompts, conversation_history, system_chart, paper, llm):
    """
    Normalize system charts for visualization (Axiomatic design version, 5 steps).

    Step 1: the model groups the structures of the chart into constituents and
            returns JSON; the code tags every Structure cell with its group(s)
            via tag_table_with_categories() (no LLM tagging step).
    Steps 2..5: the current chart is appended to the prompt and the answer is
            the chart for the next step.

    Args:
        count_for_visualization: Current normalization step number (0-based)
        prompts: List of prompts for normalization
        conversation_history: Conversation history with LLM
        system_chart: The current (tagged / normalized) system chart
        paper: Full content of the paper
        llm: Language model for processing

    Returns:
        str: Normalized system chart content
    """
    step_descriptions = {
        0: "categorizing and tagging structures",
        1: "standardizing properties",
        2: "splitting mechanisms",
        3: "separating multi-properties",
        4: "filling missing values"
    }

    step_name = step_descriptions.get(count_for_visualization, f"step {count_for_visualization}")

    print(f"\n🔄 NORMALIZATION STEP {count_for_visualization + 1}: {step_name}")

    try:
        # Handle the first step (categorization) with special processing
        if count_for_visualization == 0:
            system_chart = _clean_table_text(system_chart)
            microstructures = _parse_microstructures(system_chart)
            print(f"🔍 Extracted microstructures: {microstructures}")

            # Question 1: Initial categorization (JSON answer)
            conversation_history.append(HumanMessage(content=prompts[count_for_visualization] + microstructures + '\nOriginal paper: ' + paper))
            response = await llm.ainvoke(conversation_history)
            conversation_history.append(AIMessage(content=response.content))
            print('--------------------------------')
            print(response.content)
            print('--------------------------------')

            try:
                tagged_chart = tag_table_with_categories(system_chart, str(response.content))
            except Exception as e:
                print(f"Tagging failed ({e}); keeping the untagged chart")
                return system_chart
            return tagged_chart
        else:
            conversation_history.append(HumanMessage(content=prompts[count_for_visualization] + system_chart))
            response = await llm.ainvoke(conversation_history)
            conversation_history.append(AIMessage(content=response.content))
            cleaned = _clean_table_text(response.content)
            return cleaned if cleaned else response.content

    except Exception as e:
        print(f"❌ Error in normalization step {count_for_visualization}: {str(e)}")
        print(f"📤 Returning original chart due to error\n")
        return system_chart  # Return original chart if normalization fails

async def create_visualization(normalized_chart, paper_title, temp_dir, created_visualizations):
    """
    Create visualization from normalized system chart
    
    Args:
        normalized_chart: The normalized system chart content
        paper_title: Title of the paper
        temp_dir: Temporary directory to save visualization files
        created_visualizations: List to store created visualization file paths
    
    Returns:
        str: Path to the created visualization file
    """
    try:
        # Create the network diagram with the normalized result
        visualization_path = create_network_diagram(normalized_chart, paper_title, temp_dir)
        
        # Add to the list of created visualizations
        created_visualizations.append(visualization_path)
        
        print(f"✅ Visualization created for '{paper_title}': {visualization_path}")
        return visualization_path
        
    except Exception as e:
        print(f"❌ Error creating visualization for '{paper_title}': {str(e)}")
        return None