'''
Created on 24-Jul-2025

@author: bthakur
'''

import inspect
import traceback
import random
from IPython.display import display, HTML
from .._base import BaseQGen
import sympy as sp

class BaseSubtractQGen(BaseQGen):
    """
    Base class for subtraction question generators (decimals or integers).
    All math output is formatted with SymPy's latex() for consistency.
    Uses only SymPy number types for easy conversion to LaTeX and symbolic math.
    """

    def __init__(self, min_value, max_value, question_types=None, **kwargs):
        super().__init__(**kwargs)
        self.min_value = min_value
        self.max_value = max_value
        self.question_types = question_types or self._get_default_question_types()
        self._reset_question_queue()
        self._swp_templates = None


    def get_word_problem_templates(self):
        return self.get_swp_templates().get_templates()

    def get_swp_templates(self):
        if self._swp_templates is None:
            self._swp_templates = self._init_swp_templates()
            pass
        return self._swp_templates
        pass

    def _init_swp_templates(self):
        raise NotImplementedError
        pass

    def get_qtypes(self):
        return self.question_types

    def _get_default_question_types(self):
        return [
            'mcq_plain', 'fill_in_blank', 'true_false', 'short_answer', 'matching',
            'sorting', 'word_problem', 'calculation_working', 'explain_your_answer',
            'diagram_based', 'estimation', 'table_completion', 'pattern_recognition',
            'direct_calculation', 'explanation', 'error_analysis'
        ]

    def align_and_subtract_block(self, a, b, **kwargs):
        return self._align_and_subtract_block(a, b, **kwargs)
        pass

    def qtype_generator(self):
        core_types = [q for q in self.question_types if q != 'word_problem']
        while True:
            cycle_types = core_types[:]
            random.shuffle(cycle_types)
            interleaved = []
            for qtype in cycle_types:
                interleaved.append(qtype)
                interleaved.append('word_problem')
            for qtype in interleaved:
                yield qtype

    def _reset_question_queue(self):
        self._question_queue = self.question_types[:]
        random.shuffle(self._question_queue)

    def get_question(self, qtype=None, use_gen=False, word_template_name = None):
        if qtype is None:
            if use_gen:
                if not hasattr(self, '_qtype_gen'):
                    self._qtype_gen = self.qtype_generator()
                qtype = next(self._qtype_gen)
            else:
                if not self._question_queue:
                    self._reset_question_queue()
                qtype = self._question_queue.pop()



        print(f"Calling method qtype: {qtype}")
        method = getattr(self, qtype)

        if qtype == "word_problem":
            return method(word_template_name = word_template_name)

        return method()
        pass




    def show_question(self, q_dict, show_solution=True):
        try:
            display(HTML(f"<b>Question:</b> {q_dict['question']}"))
            opts = q_dict.get('options', None)
            if opts is not None:
                for opt in opts:
                    display(HTML(opt))
            display(HTML(f"<br><b>Answer:</b> {q_dict['answer']}"))
            if show_solution:
                display(HTML(f"<br><b>Solution:</b><br>{q_dict['solution']}"))
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}"""
            self._log(sError)
            self._log(traceback.format_exc())
            return None

    # --- Methods to override in subclasses ---
    def _rand_pair(self):
        raise NotImplementedError




    def mcq_plain(self):
        a, b = self._rand_pair()
        correct = self._result(a, b)
        distractors = set()
        while len(distractors) < 3:
            d = self._distractor(correct)
            if d != correct:
                distractors.add(d)
        options = [correct] + list(distractors)
        random.shuffle(options)
        option_letters = "ABCD"
        options_dict = {option_letters[i]: options[i] for i in range(4)}
        correct_letter = [k for k, v in options_dict.items() if v == correct][0]
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        solution = (
            f"We need to subtract <b>${self.sp_latex(b)}$</b> from <b>${self.sp_latex(a)}$</b>.<br>"
            f"Step 1: Write the numbers for subtraction.<br>"
            f"Step 2: Align them and subtract: {subtraction_block}"
            f"Step 3: The result is <b>${result_str}$</b>.<br>"
        )
        return {
            "type": "mcq_plain",
            "question": f"What is the difference between <b>${self.sp_latex(a)}$</b> and <b>${self.sp_latex(b)}$</b>?",
            "options": [f"{k}. ${self.sp_latex(v)}$" for k, v in options_dict.items()],
            "answer": f"{correct_letter}. ${self.sp_latex(correct)}$",
            "solution": solution
        }

    def fill_in_blank(self):
        a, b = self._rand_pair()
        blank = self._result(a, b)
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, blank)
        solution = (
            f"Let the blank be $x$.<br>"
            f"${self.sp_latex(a)} − x = {self.sp_latex(blank)}$<br>"
            f"<br>Change sign of each term both sides.<br>"
            f"$- {self.sp_latex(a)} + x = - {self.sp_latex(blank)}$<br>"
            f"<br>Add ${self.sp_latex(a)}$ to both sides.<br>"
            f"$\\implies + {self.sp_latex(a)} - {self.sp_latex(a)} − x =  {self.sp_latex(a)} - {self.sp_latex(blank)}$<br>"
            f"$\\implies x = {self.sp_latex(b)}$<br>"
            f"Align and subtract as shown: {subtraction_block}"
        )
        return {
            "type": "fill_in_blank",
            "question": f"${self.sp_latex(a)} −$ ___ $= {self.sp_latex(blank)}$",
            "answer": f"${self.sp_latex(b)}$",
            "solution": solution
        }

    def true_false(self):
        a, b = self._rand_pair()
        correct = self._result(a, b)
        is_true = random.choice([True, False])
        if is_true:
            statement = f"${self.sp_latex(a)} − {self.sp_latex(b)} = {self.sp_latex(correct)}$"
            answer = "True"
            a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
            solution = (
                f"Let's check the statement:<br><br>"
                f"Step 1: Align and subtract:<br>{subtraction_block}<br>"
                f"Step 2: ${self.sp_latex(a)} − {self.sp_latex(b)} = {result_str}$.<br><br>"
                f"Yes, ${self.sp_latex(a)} − {self.sp_latex(b)} = {self.sp_latex(correct)}$ is correct."
            )
        else:
            fake = correct + sp.Rational(random.choice([-2, -1, 1, 2]))
            while fake == correct:
                fake = correct + sp.Rational(random.choice([-2, -1, 1, 2]))



            statement = f"${self.sp_latex(a)} − {self.sp_latex(b)} = {self.sp_latex(fake)}$"
            answer = "False"
            a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
            solution = (
                f"Let's check the statement:<br><br>"
                f"Step 1: Align and subtract:<br>{subtraction_block}<br>"
                f"Step 2: ${self.sp_latex(a)} − {self.sp_latex(b)} = {result_str}$.<br><br>"
                f"No, ${self.sp_latex(a)} − {self.sp_latex(b)} = {self.sp_latex(fake)}$ is not correct. The correct answer is ${result_str}$."
            )
        return {
            "type": "true_false",
            "question": f"True or False: {statement}",
            "answer": answer,
            "solution": solution
        }

    def short_answer(self):
        a, b = self._rand_pair()
        result = self._result(a, b)
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        solution = (
            f"Step 1: Align and subtract:<br>{subtraction_block}<br>"
            f"Step 2: The answer is <b>${result_str}$</b>."
        )
        return {
            "type": "short_answer",
            "question": f"What is the result when you subtract ${self.sp_latex(b)}$ from ${self.sp_latex(a)}$?",
            "answer": f"${self.sp_latex(result)}$",
            "solution": solution
        }

    def matching(self):
        pairs = [self._rand_pair() for _ in range(4)]
        expressions = [f"${self.sp_latex(a)} − {self.sp_latex(b)}$" for a, b in pairs]
        answers = [self._result(a, b) for a, b in pairs]
        shuffled_answers = answers[:]
        random.shuffle(shuffled_answers)
        solutions_html = ""
        for (a, b), ans in zip(pairs, answers):
            a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
            solutions_html += f"<b>${self.sp_latex(a)} − {self.sp_latex(b)}$</b>:<br>{subtraction_block}<br>"

        roman_numerals = ['i', 'ii', 'iii', 'iv']

        # HTML table for matching with left alignment and extra space between columns
        question = (
            "Match each subtraction expression (Column A) to its correct answer (Column B):<br>"
            "<table border='1' cellpadding='10' style='border-collapse:collapse;'>"
            "<tr><th style='text-align:left;padding-right:36px;'>Column A</th>"
            "<th style='text-align:left;padding-left:24px;'>Column B</th></tr>"
        )
        for i in range(4):
            expr = f"{chr(65+i)}. {expressions[i]}"
            ans = f"{roman_numerals[i]}. ${self.sp_latex(shuffled_answers[i])}$"
            question += (
                f"<tr>"
                f"<td style='text-align:left;padding-right:36px;'>{expr}</td>"
                f"<td style='text-align:left;padding-left:24px;'>{ans}</td>"
                f"</tr>"
            )
        question += "</table>"

        answer = ", ".join([f"{chr(65+i)}-{roman_numerals[shuffled_answers.index(ans)]}" for i, ans in enumerate(answers)])
        solution = f"Correct matching (with working):<br>{solutions_html}"
        return {
            "type": "matching",
            "question": question,
            "answer": answer,
            "solution": solution
        }

    def sorting(self):
        numbers = sorted([self._rand_pair()[0] for _ in range(3)], reverse=True)
        sub = self._rand_pair()[1]
        results = [(n, self._result(n, sub)) for n in numbers]
        sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
        steps = ""
        for n, res in results:
            n_str, sub_str, res_str, subtraction_block = self.align_and_subtract_block(n, sub)
            steps += subtraction_block + "<br>"
        answer = "<br>".join([f"${self.sp_latex(x[0])} − {self.sp_latex(sub)} = {self.sp_latex(x[1])}$" for x in sorted_results])
        solution = (
            f"For each number, subtract <b>${self.sp_latex(sub)}$</b>:<br>"
            f"{steps}"
            f"Arrange results in descending order:<br>"
            f"{' > '.join([f'${self.sp_latex(x[1])}$' for x in sorted_results])}<br>"
        )
        question = (
            f"Subtract <b>${self.sp_latex(sub)}$</b> from each: "
            + ', '.join([f"${self.sp_latex(n)}$" for n in numbers])
            + ". Arrange the results in descending order."
        )
        return {
            "type": "sorting",
            "question": question,
            "answer": answer,
            "solution": solution
        }

    def word_problem(self, word_template_name = None):

        if word_template_name is None:
            template_func = random.choice(self.get_word_problem_templates())
        else:
            template_funcs = [x for x in self.get_word_problem_templates()
                              if word_template_name in x.__name__]
            template_func = random.choice(template_funcs)

        print(f""" template_func.__name__ = {template_func.__name__}  """)

        a, b = self._rand_pair()
        question, explanation = template_func(a, b)
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b,
                                                                with_steps = False)
        solution = (
            f"{explanation}<br>"
            f"{subtraction_block}<br>"
        )
        return {
            "type": "word_problem",
            "question": question,
            "answer": f"${result_str}$",
            "solution": solution
        }

    def calculation_working(self):
        a, b = self._rand_pair()
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        return {
            "type": "calculation_working",
            "question": f"Show your working to solve: ${self.sp_latex(a)} − {self.sp_latex(b)}$",
            "answer": f"${result_str}$",
            "solution": f"Column subtraction:<br>{subtraction_block}"
        }

    def explain_your_answer(self):
        a, b = self._rand_pair()
        result = self._result(a, b)
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        return {
            "type": "explain_your_answer",
            "question": f"Explain why ${self.sp_latex(a)} − {self.sp_latex(b)} = {self.sp_latex(result)}$.",
            "answer": f"Because when you subtract ${self.sp_latex(b)}$ from ${self.sp_latex(a)}$, you get ${self.sp_latex(result)}$.",
            "solution": f"Align the numbers and subtract:<br>{subtraction_block}"
        }

    def diagram_based(self):
        a, b = self._rand_pair()
        result = self._result(a, b)
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        solution = (
            f"The distance on the number line between ${self.sp_latex(a)}$ and ${self.sp_latex(b)}$ is $|{self.sp_latex(a)} − {self.sp_latex(b)}| = {self.sp_latex(abs(a-b))}$.<br>"
            f"Vertical subtraction:<br>{subtraction_block}"
        )
        return {
            "type": "diagram_based",
            "question": f"On a number line, what is the distance between ${self.sp_latex(a)}$ and ${self.sp_latex(b)}$?",
            "answer": f"${self.sp_latex(abs(a-b))}$",
            "solution": solution
        }

    def estimation(self):
        a, b = self._rand_pair()
        est_a = round(float(a))  # Cast to float for proper rounding if symbolic
        est_b = round(float(b))
        estimate = est_a - est_b

        max_len = max(len(str(est_a)), len(str(est_b)), len(str(estimate)))
        a_str = str(est_a).zfill(max_len) + '.'
        b_str = str(est_b).zfill(max_len) + '.'
        result_str = str(estimate).zfill(max_len) + '.'

        width = max_len + 4
        solid_line = "─" * width

        subtraction_block = (
            f"<pre style='font-size:1.2em; margin-bottom:0;'>"
            f"{a_str.rjust(width)}\n"
            f"{'− ' + b_str.rjust(width-2)}\n"
            f"{solid_line}\n"
            f"{result_str.rjust(width)}\n"
            f"</pre>"
        )

        solution = (
            f"Step 1: Round each number to the nearest whole number.<br>"
            f"${self.sp_latex(a)}$ rounds to <b>${self.sp_latex(est_a)}$</b>.<br>"
            f"${self.sp_latex(b)}$ rounds to <b>${self.sp_latex(est_b)}$</b>.<br>"
            f"Step 2: Subtract the rounded numbers: {subtraction_block}"
            f"Step 3: The estimated answer is <b>${self.sp_latex(estimate)}$</b>.<br>"
        )
        question = f"Estimate the result of <b>${self.sp_latex(a)}$</b> − <b>${self.sp_latex(b)}$</b>."
        return {
            "type": "estimation",
            "question": question,
            "answer": f"${self.sp_latex(estimate)}$",
            "solution": solution
        }

    def table_completion(self):
        rows = [self._rand_pair() for _ in range(3)]
        table_html = (
            "<table border='1' cellpadding='12' style='border-collapse:collapse;'>"
            "<tr>"
            "<th style='text-align:left;padding-right:28px;'>a</th>"
            "<th style='text-align:left;padding-right:28px;'>b</th>"
            "<th style='text-align:left;padding-right:28px;'>a − b</th>"
            "</tr>"
        )
        for a, b in rows:
            table_html += (
                f"<tr>"
                f"<td style='text-align:left;padding-right:28px;'>${self.sp_latex(a)}$</td>"
                f"<td style='text-align:left;padding-right:28px;'>${self.sp_latex(b)}$</td>"
                f"<td style='text-align:left;padding-right:28px;'>___</td>"
                f"</tr>"
            )
        table_html += "</table>"

        solution = ""
        for a, b in rows:
            a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
            solution += f"<b>${self.sp_latex(a)} − {self.sp_latex(b)}$</b>:<br>{subtraction_block}<br>"
        answer = ", ".join([f"${self.sp_latex(self._result(a, b))}$" for a, b in rows])
        return {
            "type": "table_completion",
            "question": f"Complete the table:<br>{table_html}",
            "answer": answer,
            "solution": solution
        }



    def pattern_recognition(self):
        start, diff = self._rand_pair()
        # Ensure diff is not zero
        while abs(diff) < 1e-6 if isinstance(diff, float) else diff == 0:
            start, diff = self._rand_pair()
        # Generate the pattern (using addition for progression)
        pattern = [start + i * diff for i in range(4)]
        next_term = pattern[-1] + diff
        # Calculate actual differences and get subtraction blocks
        differences = []
        subtraction_blocks = []
        for i in range(1, len(pattern)):
            a, b = pattern[i], pattern[i-1]
            diff_val = a - b
            differences.append(diff_val)
            _, _, _, block = self.align_and_subtract_block(a, b, with_steps = False)
            subtraction_blocks.append(block)
        # Block for next term
        _, _, _, next_term_block = self.align_and_subtract_block(pattern[-1], differences[0],
                                                                 with_steps = False,
                                                                 op = "+")
        # Build solution
        solution = "<b>Step 1: Find the difference between consecutive terms using column subtraction:</b><br>"
        for i in range(1, len(pattern)):
            solution += (
                f"${self.sp_latex(pattern[i-1])} \\to {self.sp_latex(pattern[i])}$: "
                f"Difference = ${self.sp_latex(differences[i-1])}$<br>"
                f"<div style='color:#2277bb;margin-left:2em'>{subtraction_blocks[i-1]}</div><br>"
            )

        solution += "<br><b>Step 2: Check if the difference is common for all pairs.</b><br>"
        all_equal = all(d == differences[0] for d in differences)
        if all_equal:
            solution += f"Yes, all differences are <b>${self.sp_latex(differences[0])}$</b>.<br>"
        else:
            solution += "The differences are not the same. This is not an arithmetic sequence.<br>"
        solution += "<br><b>Step 3: Find the next term by adding the common difference to the last term:</b><br>"
        solution += (
            f"${self.sp_latex(pattern[-1])} + {self.sp_latex(differences[0])} = {self.sp_latex(next_term)}$<br>"
            f"<div style='color:#2277bb;margin-left:2em'>{next_term_block}</div><br>"
        )
        # Prepare question
        question = (
            f"Find the next number in the pattern:<br>"
            f"{', '.join([f'${self.sp_latex(p)}$' for p in pattern])}, ___"
        )
        return {
            "type": "pattern_recognition",
            "question": question,
            "answer": f"${self.sp_latex(next_term)}$",
            "solution": solution
        }

    def direct_calculation(self):
        a, b = self._rand_pair()
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        solution = (
            f"Subtract <b>${self.sp_latex(b)}$</b> from <b>${self.sp_latex(a)}$</b>: {subtraction_block}"
            f"The answer is <b>${result_str}$</b>.<br>"
        )
        return {
            "type": "direct_calculation",
            "question": f"Subtract: ${self.sp_latex(a)} − {self.sp_latex(b)} = \\underline{{\\phantom{{0}}}}$",
            "answer": f"${result_str}$",
            "solution": solution
        }

    def explanation(self):
        a, b = self._rand_pair()
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        solution = (
            f"Step 1: Prepare both numbers for subtraction.<br>"
            f"Step 2: Align and subtract: {subtraction_block}"
            f"Step 3: The answer is <b>${result_str}$</b>.<br>"
        )
        question = f"Explain how to subtract <b>${self.sp_latex(a)}$</b> and <b>${self.sp_latex(b)}$</b>."
        answer = f"${result_str}$"
        return {
            "type": "explanation",
            "question": question,
            "answer": answer,
            "solution": solution
        }

    def error_analysis(self):
        a, b = self._rand_pair()
        correct = self._result(a, b)
        student_answer = correct if random.random() < 0.5 else correct + random.choice([-2, -1, 1, 2])
        a_str, b_str, result_str, subtraction_block = self.align_and_subtract_block(a, b)
        solution = (
            f"Step 1: Prepare the numbers for subtraction.<br>"
            f"Step 2: Subtract as shown: {subtraction_block}"
            f"Step 3: The correct answer is <b>${result_str}$</b>.<br>"
            f"Step 4: The student wrote <b>${self.sp_latex(student_answer)}$</b>.<br>"
            +
            ("The student's answer is <b>correct</b>.<br>" if student_answer == correct
             else f"The student's answer is <b>incorrect</b>.<br>The correct answer is <b>${result_str}$</b>.<br>")
        )
        question = (f"A student calculated <b>${self.sp_latex(a)} − {self.sp_latex(b)}$</b> and got <b>${self.sp_latex(student_answer)}$</b>.<br>"
                    "Is this answer correct? If not, what is the correct answer?")
        return {
            "type": "error_analysis",
            "question": question,
            "answer": f"${result_str}$",
            "solution": solution
        }

    # Helper for MCQ distractors
    def _distractor(self, correct):
        # For integers: +/-1 or +/-2; for decimals, add/subtract small float
        if isinstance(correct, sp.Integer):
            return correct + random.choice([-2, -1, 1, 2])
        elif isinstance(correct, (int, float)):
            return correct + random.choice([-2, -1, 1, 2])
        else:
            return correct + sp.Rational(random.choice([-2, -1, 1, 2]))

    def _result(self, a, b):
        return a - b






