import re
import sys


def convert_inline_math(text):
    # This pattern looks for a literal '\(' and then non-greedily
    # matches everything '.*?' until it finds a literal '\)'.
    pattern = r'\\\((.*?)\\\)'
    # The replacement adds spaces around the math expression for better rendering.
    return re.sub(pattern, r' $\1$ ', text, flags=re.DOTALL)


def convert_display_math(text):
    pattern = r'\\\[(.*?)\\\]'
    return re.sub(pattern, r'$$\1$$', text, flags=re.DOTALL)


def clean_nested_tabular(content):
    """Extract content from nested tabular environments"""

    def extract_nested_tabular(match):
        inner_content = match.group(1)
        # Split by & and \\ to get cells, join with line breaks
        lines = inner_content.split('\\\\')
        cleaned_lines = []
        for line in lines:
            cells = [cell.strip() for cell in line.split('&') if cell.strip()]
            if cells:
                cleaned_lines.extend(cells)
        return ' \\\\ '.join(cleaned_lines) if cleaned_lines else inner_content

    # Handle nested tabular environments
    while re.search(r'\\begin\{tabular\}\{[^}]*\}(.*?)\\end\{tabular\}', content, re.DOTALL):
        content = re.sub(r'\\begin\{tabular\}\{[^}]*\}(.*?)\\end\{tabular\}', extract_nested_tabular, content,
                         flags=re.DOTALL)

    return content.strip()


def parse_latex_table(table_content):
    """Convert LaTeX tabular to markdown table with support for multicolumn/multirow."""
    lines = []
    header_separator_added = False

    # Extract just the tabular content
    table_content = re.sub(r'\\begin\{tabular\}\{[^}]+\}', '', table_content)
    table_content = re.sub(r'\\end\{tabular\}', '', table_content)

    # Clean nested tabular environments first
    table_content = clean_nested_tabular(table_content)

    # Split rows by '\\' - use the original simple approach for most cases
    rows = re.split(r'\\\\', table_content)

    # Storage for multirow spans: {col_index: (remaining_span, content)}
    active_multirows = {}

    for row_idx, row in enumerate(rows):
        row = row.strip()
        if not row:
            continue

        # Remove hline/cline
        row = re.sub(r'\\hline', '', row)
        row = re.sub(r'\\cline\{[^}]+\}', '', row)
        row = row.strip()

        if not row:
            continue

        # Handle multicolumn: expand into multiple identical cells
        def expand_multicolumn(match):
            cols = int(match.group(1))
            content = match.group(2).strip()
            return ' & '.join([content] * cols)

        row = re.sub(r'\\multicolumn\{(\d+)\}\{[^}]+\}\{([^}]*)\}', expand_multicolumn, row)

        # Handle multirow: capture and store for future rows
        multirow_positions = {}

        def capture_multirow(match):
            span = int(match.group(1))
            content = match.group(2).strip()
            # We'll determine the position after splitting by &
            return f"MULTIROW:{span}:{content}"

        row = re.sub(r'\\multirow\{(\d+)\}\{[^}]*\}\{([^}]*)\}', capture_multirow, row)

        # Split into cells
        cells = [cell.strip() for cell in row.split('&')]

        # Process multirow markers and build final cells
        final_cells = []
        for col_idx, cell in enumerate(cells):
            if cell.startswith("MULTIROW:"):
                parts = cell.split(":", 2)
                span = int(parts[1])
                content = parts[2] if len(parts) > 2 else ""
                final_cells.append(content)
                # Store for future rows (span-1 because current row counts)
                active_multirows[col_idx] = (span - 1, content)
            elif col_idx in active_multirows:
                # This column has an active multirow from a previous row
                span, content = active_multirows[col_idx]
                final_cells.append(content)
                # Decrease span
                active_multirows[col_idx] = (span - 1, content)
                if span - 1 <= 0:
                    del active_multirows[col_idx]
                # Also include the original cell content if it's not empty
                if cell and cell != content:
                    final_cells[-1] = cell  # Override with current content
            else:
                final_cells.append(cell)

        # Fill any remaining active multirows for columns beyond our current cells
        max_col = max(len(final_cells), max(active_multirows.keys(), default=-1) + 1)
        while len(final_cells) < max_col:
            col_idx = len(final_cells)
            if col_idx in active_multirows:
                span, content = active_multirows[col_idx]
                final_cells.append(content)
                active_multirows[col_idx] = (span - 1, content)
                if span - 1 <= 0:
                    del active_multirows[col_idx]
            else:
                final_cells.append('')

        # Decrement remaining active multirows
        to_delete = []
        for col_idx in active_multirows:
            span, content = active_multirows[col_idx]
            active_multirows[col_idx] = (span - 1, content)
            if span - 1 <= 0:
                to_delete.append(col_idx)
        for col_idx in to_delete:
            del active_multirows[col_idx]

        # Build markdown row
        md_row = '| ' + ' | '.join(final_cells) + ' |'
        lines.append(md_row)

        # Add header separator after first row
        if not header_separator_added:
            separator = '|' + '---|' * len(final_cells)
            lines.append(separator)
            header_separator_added = True

    return '\n'.join(lines)


def parse_pipe_table(content):
    """
    Parses table content that uses pipes for columns and newlines for rows,
    which is not standard LaTeX but appears in some Mathpix outputs.
    """
    lines = content.strip().split('\n')
    md_lines = []
    header_separator_added = False

    table_data_lines = []
    caption_text_line = None

    # Separate caption from table data. Heuristic: assumes caption is the
    # last non-empty line starting with "Table X:".
    potential_caption_lines = [
        (i, line.strip()) for i, line in enumerate(lines)
        if re.match(r'^\s*Table \d+:', line.strip(), re.IGNORECASE)
    ]
    if potential_caption_lines:
        caption_index, caption_full_line = potential_caption_lines[-1]
        caption_text_line = re.sub(r'^\s*Table \d+:\s*', '', caption_full_line)
        table_data_lines = lines[:caption_index]
    else:
        table_data_lines = lines

    for line in table_data_lines:
        line = line.strip()
        # Skip empty lines and pre-existing markdown separators
        if not line or line.startswith('|---'):
            continue

        # Handle multirow/multicolumn before splitting (simplistic approach)
        line = re.sub(r'\\multirow\{[^}]+\}\{[^}]+\}\{([^}]+)\}', r'\1', line)
        line = re.sub(r'\\multicolumn\{[^}]+\}\{[^}]+\}\{([^}]+)\}', r'\1', line)

        cells = [cell.strip() for cell in line.split('|')]

        # Clean up empty cells from leading/trailing pipes
        if cells and not cells[0]:
            cells.pop(0)
        if cells and not cells[-1]:
            cells.pop()

        if not cells:
            continue

        md_row = '| ' + ' | '.join(cells) + ' |'
        md_lines.append(md_row)

        if not header_separator_added:
            separator = '|' + '---|' * len(cells)
            md_lines.append(separator)
            header_separator_added = True

    result = '\n'.join(md_lines)
    if caption_text_line:
        return f"**Table:** {caption_text_line}\n\n{result}"

    return result


def convert_tables(text):
    """Convert all LaTeX tables to markdown tables"""
    table_pattern = r'\\begin\{table\}.*?\\end\{table\}'
    tabular_pattern = r'\\begin\{tabular\}.*?\\end\{tabular\}'

    # First handle tables wrapped in the 'table' environment
    def replace_table(match):
        table_content = match.group(0)

        # Check for standard LaTeX tabular environment first
        tabular_match = re.search(tabular_pattern, table_content, re.DOTALL)
        if tabular_match:
            md_table = parse_latex_table(tabular_match.group(0))
            # Extract caption from \caption{} command
            caption_match = re.search(r'\\caption\{([^}]+)\}', table_content)
            if caption_match:
                caption = caption_match.group(1)
                return f"**Table:** {caption}\n\n{md_table}"
            return md_table
        else:
            # Fallback for non-standard, pipe-based tables
            # Clean the inner content before parsing
            inner_content = re.sub(r'\\begin\{table\}|\\end\{table\}', '', table_content)
            inner_content = re.sub(r'\\caption\{[^}]+\}|\\label\{[^}]+\}', '', inner_content)
            inner_content = re.sub(r'\\centering', '', inner_content)
            return parse_pipe_table(inner_content)

    text = re.sub(table_pattern, replace_table, text, flags=re.DOTALL)

    # Then handle any standalone 'tabular' environments
    def replace_tabular(match):
        return parse_latex_table(match.group(0))

    text = re.sub(tabular_pattern, replace_tabular, text, flags=re.DOTALL)
    return text


def remove_latex_commands(text):
    """Remove various LaTeX commands that don't have markdown equivalents"""
    # Remove footnotes
    text = re.sub(r'\\footnote\{[^}]+\}', '', text)
    # Remove some common, simple environments
    text = re.sub(r'\\begin\{(?:figure|center)\}', '', text)
    text = re.sub(r'\\end\{(?:figure|center)\}', '', text)
    # Remove smallmatrix environments (convert to simple text)
    text = re.sub(r'\\begin\{smallmatrix\}(.*?)\\end\{smallmatrix\}', r'[\1]', text, flags=re.DOTALL)
    # Remove \left and \right which are not needed in standard Markdown math
    text = re.sub(r'\\left|\\right', '', text)
    # Remove centering command
    text = re.sub(r'\\centering', '', text)
    # Clean up excessive newlines
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text


def mathpix_to_markdown(text):
    """Main conversion function applying all transformations in order."""
    # The order of operations is important.
    text = convert_display_math(text)
    text = convert_inline_math(text)
    text = convert_tables(text)
    text = remove_latex_commands(text)
    return text.strip()


def process_file(input_file, output_file):
    """Reads an input file, converts it, and saves to the output file."""
    try:
        print(f"Reading from '{input_file}'...")
        with open(input_file, 'r', encoding='utf-8') as f:
            content = f.read()

        converted_content = mathpix_to_markdown(content)

        print(f"Writing to '{output_file}'...")
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(converted_content)

        print("Conversion complete.")

    except FileNotFoundError:
        print(f"Error: The file '{input_file}' was not found.", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"An error occurred during processing: {e}", file=sys.stderr)
        sys.exit(1)


def main(in_path, out_path):
    process_file(in_path, out_path)


if __name__ == "__main__":
    main(r"/Users/arielkeslassy/PycharmProjects/PaperParser/Deep Residual Learning for Image Recognition.md",
         r"/Users/arielkeslassy/PycharmProjects/PaperParser/output.md")