# =========================================================
# Semantic Evaluation Pipeline for LaTeX Mathematical Expressions
# =========================================================

import re
import json
import os
from typing import Optional, Tuple, Dict, List

import sympy as sp
from sympy.parsing.latex import parse_latex

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel


# =========================================================
# 1. Heuristic Analyzer
# =========================================================

class HeuristicAnalyzer:
    OCR_GARBAGE_PATTERNS = [
        r"\\underline", r"\\rlap", r"\\mit", r"\\bf"
    ]

    INVALID_LATEX_PATTERNS = [
        r"\\backslash\s+[A-Za-z]",
    ]

    def check(self, formula: str) -> Dict:
        for p in self.OCR_GARBAGE_PATTERNS:
            if re.search(p, formula):
                return {
                    "syntactically_correct": False,
                    "mathematically_correct": False,
                    "semantic_verdict": "garbled",
                    "error_type": "garbled_ocr",
                    "confidence": 0.95,
                    "source": "heuristic",
                    "explanation": "The formula contains OCR or encoding artifacts.",
                    "suggested_fix": "Clean or retype the LaTeX."
                }

        for p in self.INVALID_LATEX_PATTERNS:
            if re.search(p, formula):
                return {
                    "syntactically_correct": False,
                    "mathematically_correct": False,
                    "semantic_verdict": "invalid_notation",
                    "error_type": "notation_error",
                    "confidence": 0.95,
                    "source": "heuristic",
                    "explanation": "Invalid LaTeX command usage.",
                    "suggested_fix": "Replace with valid LaTeX."
                }

        return {"pass": True}


# =========================================================
# 2. Symbolic Verifier
# =========================================================

class SymbolicVerifier:

    @staticmethod
    def normalize_latex(latex: str) -> str:
        s = latex.strip()
        s = re.sub(r"(\d)([a-zA-Z])", r"\1*\2", s)
        s = re.sub(r"([a-zA-Z])([a-zA-Z])", r"\1*\2", s)
        s = s.replace("^", "**")
        return s

    @staticmethod
    def split_equation(latex: str) -> Optional[Tuple[str, str]]:
        if latex.count("=") != 1:
            return None
        lhs, rhs = latex.split("=")
        return lhs.strip(), rhs.strip()

    def equivalent(self, lhs_latex: str, rhs_latex: str) -> Optional[bool]:
        try:
            lhs = parse_latex(lhs_latex)
            rhs = parse_latex(rhs_latex)
            return sp.simplify(lhs - rhs) == 0
        except Exception:
            pass

        try:
            lhs = sp.sympify(self.normalize_latex(lhs_latex))
            rhs = sp.sympify(self.normalize_latex(rhs_latex))
            return sp.simplify(lhs - rhs) == 0
        except Exception:
            return None


# =========================================================
# 3. Explanation-Only LLM
# =========================================================

class ExplanationLLM:
    def __init__(
        self,
        base_model: str = "Qwen/Qwen2.5-Math-1.5B-Instruct",
        lora_dir: Optional[str] = None,
        device: str = "cuda"
    ):
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(base_model)

        model = AutoModelForCausalLM.from_pretrained(
            base_model,
            device_map="auto",
            torch_dtype=torch.float16
        )

        if lora_dir:
            model = PeftModel.from_pretrained(model, lora_dir, device_map="auto")

        self.model = model

    def explain(self, formula: str, issue: str) -> str:
        prompt = f"""You are a mathematical LaTeX reviewer.

Formula:
{formula}

Issue:
{issue}

Explain:
- What is happening mathematically
- Why it cannot be verified or is incorrect
- What should be checked or fixed

Do not output True/False.
"""

        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)

        with torch.no_grad():
            output = self.model.generate(
                **inputs,
                max_new_tokens=800,
                do_sample=False,
                eos_token_id=self.tokenizer.eos_token_id
            )

        decoded = self.tokenizer.decode(output[0], skip_special_tokens=True)
        return decoded.replace(prompt, "").strip()


# =========================================================
# 4. Formula Evaluator
# =========================================================

class FormulaEvaluator:
    def __init__(self, explainer: ExplanationLLM):
        self.heuristic = HeuristicAnalyzer()
        self.symbolic = SymbolicVerifier()
        self.explainer = explainer

    def evaluate(self, formula: str) -> Dict:
        h = self.heuristic.check(formula)
        if "pass" not in h:
            return h

        eq = self.symbolic.split_equation(formula)
        if eq:
            result = self.symbolic.equivalent(eq[0], eq[1])

            if result is True:
                return {
                    "syntactically_correct": True,
                    "mathematically_correct": True,
                    "semantic_verdict": "equivalent",
                    "error_type": None,
                    "confidence": 0.99,
                    "source": "sympy",
                    "explanation": "Both sides are mathematically equivalent.",
                    "suggested_fix": None
                }

            if result is False:
                return {
                    "syntactically_correct": True,
                    "mathematically_correct": False,
                    "semantic_verdict": "contradiction",
                    "error_type": "math_error",
                    "confidence": 0.99,
                    "source": "sympy+llm",
                    "explanation": self.explainer.explain(
                        formula, "Equality does not hold in general"
                    ),
                    "suggested_fix": "Correct the mathematical identity."
                }

        return {
            "syntactically_correct": True,
            "mathematically_correct": None,
            "semantic_verdict": "non_verifiable",
            "error_type": "uncertain",
            "confidence": 0.5,
            "source": "llm",
            "explanation": self.explainer.explain(
                formula, "Expression is not a verifiable equality"
            ),
            "suggested_fix": "Manual inspection required."
        }


# =========================================================
# 5. LaTeX Processor
# =========================================================

class LatexProcessor:
    MATH_ENV_PATTERNS = [
        r"\$(.*?)\$",
        r"\\\((.*?)\\\)",
        r"\\\[(.*?)\\\]",
        r"\\begin\{equation\}(.*?)\\end\{equation\}",
    ]

    def extract(self, tex: str) -> List[str]:
        formulas = []
        for p in self.MATH_ENV_PATTERNS:
            formulas.extend(re.findall(p, tex, re.DOTALL))
        return [f.strip() for f in formulas]

    def clean(self, formula: str) -> str:
        formula = re.sub(r"\\underline\{\{.*?\}\}", "", formula)
        formula = re.sub(r"\\mit|\\bf", "", formula)
        formula = re.sub(r"\\l\\ell|\\L", "", formula)
        formula = formula.replace(r"\rlap/", "")
        formula = re.sub(r"[\^\_]{2,}", "", formula)
        formula = re.sub(r"\{\s*\}", "", formula)
        return formula.strip()

    def fix(self, formula: str) -> str:
        formula = re.sub(r"\\backslash\s*([A-Za-z])", r"\\setminus \1", formula)
        formula = formula.replace(r"\implies", r"\Rightarrow")
        return formula


# =========================================================
# 6. High-Level Pipeline
# =========================================================

class SemanticEvaluationPipeline:
    def __init__(
        self,
        base_model: str = "Qwen/Qwen2.5-Math-1.5B-Instruct",
        lora_dir: Optional[str] = None,
        device: str = "cuda"
    ):
        self.latex = LatexProcessor()
        self.explainer = ExplanationLLM(base_model, lora_dir, device)
        self.evaluator = FormulaEvaluator(self.explainer)

    def evaluate_formula(self, formula: str) -> Dict:
        cleaned = self.latex.clean(formula)
        fixed = self.latex.fix(cleaned)
        return {
            "original_formula": formula,
            "fixed_formula": fixed,
            **self.evaluator.evaluate(fixed)
        }

    def evaluate_tex_file(self, tex_path: str) -> List[Dict]:
        with open(tex_path, "r", encoding="utf-8") as f:
            tex = f.read()

        formulas = self.latex.extract(tex)
        return [self.evaluate_formula(f) for f in formulas]

    def save_results(self, results: List[Dict], out_dir: str):
        os.makedirs(out_dir, exist_ok=True)

        # JSON
        json_path = os.path.join(out_dir, "results_test_1.json")
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)

        # TXT
        txt_path = os.path.join(out_dir, "results_test_1.txt")
        with open(txt_path, "w", encoding="utf-8") as f:
            for i, r in enumerate(results, 1):
                f.write(f"Formula #{i}\n")
                f.write("-" * 80 + "\n")
                f.write(f"Original Formula:\n{r['original_formula']}\n\n")
                f.write(f"Fixed Formula:\n{r['fixed_formula']}\n\n")
                f.write(f"Syntactically Correct: {r['syntactically_correct']}\n")
                f.write(f"Mathematically Correct: {r['mathematically_correct']}\n")
                f.write(f"Semantic Verdict: {r['semantic_verdict']}\n")
                f.write(f"Error Type: {r['error_type']}\n")
                f.write(f"Confidence: {r['confidence']}\n")
                f.write(f"Source: {r['source']}\n\n")
                f.write("Explanation:\n")
                f.write(r["explanation"] + "\n")
                f.write("=" * 80 + "\n\n")

        return json_path, txt_path


# =========================================================
# 7. CLI / Main
# =========================================================

def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Semantic evaluation of LaTeX mathematical expressions"
    )
    parser.add_argument("tex_file", help="Path to LaTeX (.tex) file")
    parser.add_argument("--out", default="output", help="Output directory")
    parser.add_argument("--base_model", default="Qwen/Qwen2.5-Math-1.5B-Instruct")
    parser.add_argument("--lora_dir", default=None)
    parser.add_argument("--device", default="cuda")

    args = parser.parse_args()

    pipeline = SemanticEvaluationPipeline(
        base_model=args.base_model,
        lora_dir=args.lora_dir,
        device=args.device
    )

    results = pipeline.evaluate_tex_file(args.tex_file)
    json_path, txt_path = pipeline.save_results(results, args.out)

    print(f"✓ Evaluated {len(results)} formulas")
    print(f"✓ JSON saved to: {json_path}")
    print(f"✓ TXT  saved to: {txt_path}")


if __name__ == "__main__":
    main()