from abc import ABC, abstractmethod
from typing import Any, Generator

from pysat.examples.lbx import LBX as MCSSolver # type: ignore
from pysat.examples.musx import MUSX as MUSSolver # type: ignore
from pysat.examples.rc2 import RC2 as MaxSATSolver # type: ignore
from pysat.solvers import Solver # type: ignore
from pysat.examples.lsu import LSU # type: ignore
from pysat.examples.hitman import Hitman # type: ignore

from satquest.cnf import CNF
from satquest.constants import GIT_HASH, SAT_SOLVER_NAME
from satquest.question import Question
from satquest.satquest_utils import cnf2wcnf, get_class_source_hash


class Problem(ABC):
    def __init__(self, cnf: CNF):
        self.cnf = cnf
        self._solution = None
        self._solver_metadata = None

    @property
    @abstractmethod
    def solution(self) -> Any | None:
        pass

    @abstractmethod
    def solution_enumerate(self) -> Generator[str | None, None, None]:
        pass

    @property
    @abstractmethod
    def search_space_size(self) -> Any:
        pass

    @abstractmethod
    def check(self, answer: Any) -> bool:
        pass

    @abstractmethod
    def format_check(self, answer: Any) -> bool:
        pass

    @abstractmethod
    def accept(self, question: Question, *args, **kwargs) -> str:
        pass

    @property
    @abstractmethod
    def ANSWER_PATTERN(self) -> str:
        pass

    @property
    def solver_metadata(self) -> dict | None:
        if self._solver_metadata is None:
            _ = self.solution
        return self._solver_metadata

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}_{GIT_HASH}_{get_class_source_hash(self.__class__)}"


class SATDP(Problem):
    @property
    def solution(self) -> str | None:
        if self._solution is None:
            self._solution, self._solver_metadata = None, None
            try:
                with Solver(name=SAT_SOLVER_NAME, bootstrap_with=self.cnf.clauses) as solver:
                    self._solution = str(int(solver.solve()))
                    self._solver_metadata = {**solver.accum_stats(), "solvers": (SAT_SOLVER_NAME,)}
            except Exception:
                pass
        return self._solution

    def solution_enumerate(self) -> Generator[str | None, None, None]:
        yield self.solution

    @property
    def search_space_size(self) -> Any:
        return 2**self.cnf.nv

    def check(self, answer: str) -> bool:
        try:
            assert self.format_check(answer)
            return answer == str(self.solution)
        except Exception:
            pass
        return False

    def format_check(self, answer: str) -> bool:
        return isinstance(answer, str) and len(answer) == 1 and set(answer).issubset({"0", "1"})

    def accept(self, question: Question, *args, **kwargs) -> str:
        return question.visit_satdp(self.cnf, *args, **kwargs)

    @property
    def ANSWER_PATTERN(self) -> str:
        return r"(?=([01]{%d}))" % 1


class SATSP(Problem):
    @property
    def solution(self) -> str | None:
        assert self.cnf.is_sat
        if self._solution is None:
            self._solution, self._solver_metadata = None, None
            try:
                with Solver(name=SAT_SOLVER_NAME, bootstrap_with=self.cnf.clauses) as solver:
                    _is_sat = solver.solve()
                    if _is_sat:
                        self._solution = "".join(["1" if iv > 0 else "0" for iv in solver.get_model()])
                    self._solver_metadata = {**solver.accum_stats(), "solvers": (SAT_SOLVER_NAME,)}
            except Exception:
                pass
        return self._solution

    def solution_enumerate(self) -> Generator[str, None, None]:
        assert self.cnf.is_sat
        with Solver(name=SAT_SOLVER_NAME, bootstrap_with=self.cnf.clauses) as solver:
            for s in solver.enum_models():
                yield "".join(["1" if iv > 0 else "0" for iv in s])

    @property
    def search_space_size(self) -> Any:
        return 2**self.cnf.nv

    def check(self, answer: str) -> bool:
        assert self.cnf.is_sat
        try:
            assert self.format_check(answer)
            with Solver(name=SAT_SOLVER_NAME, bootstrap_with=self.cnf.clauses) as solver:
                return solver.solve(assumptions=[(i + 1) if ai == "1" else -(i + 1) for i, ai in enumerate(answer)])
        except Exception:
            pass
        return False

    def format_check(self, answer: str) -> bool:
        return isinstance(answer, str) and len(answer) == self.cnf.nv and set(answer).issubset({"0", "1"})

    def accept(self, question: Question, *args, **kwargs) -> str:
        return question.visit_satsp(self.cnf, *args, **kwargs)

    @property
    def ANSWER_PATTERN(self) -> str:
        return r"(?=([01]{%d}))" % self.cnf.nv


class MaxSAT(Problem):
    @property
    def solution(self) -> str | None:
        assert not self.cnf.is_sat
        if self._solution is None:
            self._solution, self._solver_metadata = None, None
            try:
                with MaxSATSolver(cnf2wcnf(self.cnf.cnf), solver=SAT_SOLVER_NAME, verbose=0) as solver:
                    solver.compute()
                    self._solver_metadata = {**solver.oracle.accum_stats(), "solvers": (SAT_SOLVER_NAME, "RC2")}
                    self._solution = "".join(["1" if iv > 0 else "0" for iv in solver.model])
            except Exception:
                pass
        return self._solution

    def solution_enumerate(self) -> Generator[str, None, None]:
        assert not self.cnf.is_sat

        with MaxSATSolver(cnf2wcnf(self.cnf.cnf), solver=SAT_SOLVER_NAME, verbose=0) as solver:
            pre_cost = None
            for s in solver.enumerate():
                if pre_cost is None:
                    pre_cost = solver.cost
                if solver.cost != pre_cost:
                    break
                yield "".join(["1" if iv > 0 else "0" for iv in s])

    @property
    def search_space_size(self) -> Any:
        return 2**self.cnf.nv

    def check(self, answer: str) -> bool:
        assert not self.cnf.is_sat
        try:
            assert self.format_check(answer)
            with LSU(cnf2wcnf(self.cnf.cnf), solver=SAT_SOLVER_NAME, verbose=0) as solver:
                solver.solve()
                answer_cost = solver._get_model_cost(
                    cnf2wcnf(self.cnf.cnf), [(i + 1) if ai == "1" else -(i + 1) for i, ai in enumerate(answer)]
                )
                return answer_cost == solver.cost
        except Exception:
            pass
        return False

    def format_check(self, answer: str) -> bool:
        return isinstance(answer, str) and len(answer) == self.cnf.nv and set(answer).issubset({"0", "1"})

    def accept(self, question: Question, *args, **kwargs) -> str:
        return question.visit_maxsat(self.cnf, *args, **kwargs)

    @property
    def ANSWER_PATTERN(self) -> str:
        return r"(?=([01]{%d}))" % self.cnf.nv


class MCS(Problem):
    @property
    def solution(self) -> str | None:
        assert not self.cnf.is_sat
        if self._solution is None:
            with MCSSolver(cnf2wcnf(self.cnf.cnf), use_cld=False, solver_name=SAT_SOLVER_NAME) as solver:
                _solution_model = solver.compute()
                self._solution = "".join(["1" if i in _solution_model else "0" for i in range(1, self.cnf.mc + 1)])
                self._solver_metadata = {**solver.oracle.accum_stats(), "solvers": (SAT_SOLVER_NAME, "LBX")}
        return self._solution

    def solution_enumerate(self) -> Generator[str, None, None]:
        assert not self.cnf.is_sat
        with MCSSolver(cnf2wcnf(self.cnf.cnf), use_cld=False, solver_name=SAT_SOLVER_NAME) as solver:
            for mcs in solver.enumerate():
                solver.block(mcs)
                yield "".join(["1" if i in mcs else "0" for i in range(1, len(self.cnf.clauses) + 1)])

    @property
    def search_space_size(self) -> Any:
        return 2**self.cnf.mc

    def check(self, answer: str) -> bool:
        assert not self.cnf.is_sat
        try:
            assert self.format_check(answer)
            unmcs_clauses = [self.cnf.clauses[i] for i in range(self.cnf.mc) if answer[i] == "0"]
            mcs_clauses = [self.cnf.clauses[i] for i in range(self.cnf.mc) if answer[i] == "1"]
            with Solver(name=SAT_SOLVER_NAME, bootstrap_with=CNF(unmcs_clauses).cnf) as solver:
                if not solver.solve():
                    return False
            for i in range(len(mcs_clauses)):
                add_mcs_clauses = unmcs_clauses + [mcs_clauses[i]]
                with Solver(name=SAT_SOLVER_NAME, bootstrap_with=CNF(add_mcs_clauses).cnf) as solver:
                    if solver.solve():
                        return False
            return True
        except Exception:
            pass
        return False

    def format_check(self, answer: str) -> bool:
        return isinstance(answer, str) and len(answer) == self.cnf.mc and set(answer).issubset({"0", "1"})

    def accept(self, question: Question, *args, **kwargs) -> str:
        return question.visit_mcs(self.cnf, *args, **kwargs)

    @property
    def ANSWER_PATTERN(self) -> str:
        return r"(?=([01]{%d}))" % self.cnf.mc


class MUS(Problem):
    @property
    def solution(self) -> str | None:
        assert not self.cnf.is_sat
        if self._solution is None:
            with MUSSolver(self.cnf.cnf, solver=SAT_SOLVER_NAME, verbosity=0) as solver:
                _solution_model = solver.compute()
                self._solution = "".join(["1" if i in _solution_model else "0" for i in range(1, self.cnf.mc + 1)])
                self._solver_metadata = {**solver.oracle.accum_stats(), "solvers": (SAT_SOLVER_NAME, "MUSX")}
        return self._solution

    def solution_enumerate(self) -> Generator[str, None, None]:
        assert not self.cnf.is_sat

        with Hitman(solver="m22") as hitman_solver:
            with MCSSolver(cnf2wcnf(self.cnf.cnf), use_cld=False, solver_name=SAT_SOLVER_NAME) as solver:
                for mcs in solver.enumerate():
                    solver.block(mcs)
                    hitman_solver.hit(mcs)
            for mus in hitman_solver.enumerate():
                yield "".join(["1" if i in mus else "0" for i in range(1, len(self.cnf.clauses) + 1)])

    @property
    def search_space_size(self) -> Any:
        return 2**self.cnf.mc

    def check(self, answer: str) -> bool:
        assert not self.cnf.is_sat
        try:
            assert self.format_check(answer)
            mus_clauses = [self.cnf.clauses[i] for i in range(self.cnf.mc) if answer[i] == "1"]
            with Solver(name=SAT_SOLVER_NAME, bootstrap_with=CNF(mus_clauses).cnf) as solver:
                if solver.solve():
                    return False
            for i in range(len(mus_clauses)):
                sub_mus_clauses = mus_clauses[:i] + mus_clauses[i + 1 :]
                with Solver(name=SAT_SOLVER_NAME, bootstrap_with=CNF(sub_mus_clauses).cnf) as solver:
                    if not solver.solve():
                        return False
            return True
        except Exception:
            pass
        return False

    def format_check(self, answer: str) -> bool:
        return isinstance(answer, str) and len(answer) == self.cnf.mc and set(answer).issubset({"0", "1"})

    def accept(self, question: Question, *args, **kwargs) -> str:
        return question.visit_mus(self.cnf, *args, **kwargs)

    @property
    def ANSWER_PATTERN(self) -> str:
        return r"(?=([01]{%d}))" % self.cnf.mc


def create_problem(problem_type: str, cnf: CNF) -> Problem:
    problem_type = problem_type.lower()
    match problem_type:
        case "satdp" | "satdp_sat" | "satdp_unsat":
            return SATDP(cnf)
        case "satsp":
            return SATSP(cnf)
        case "maxsat":
            return MaxSAT(cnf)
        case "mcs":
            return MCS(cnf)
        case "mus":
            return MUS(cnf)
    raise ValueError(f"Invalid problem type: {problem_type}")