import queue
from dataclasses import dataclass
from functools import cached_property
from itertools import product

import z3

from common import config
from common.parallelism import JobResult, Pool
from common.relations import Relation
from common.z3_utils.expr_utils import uninterpreted_sorts
from common.z3_utils.pickling import PicklableExpr, PicklableSort
from fest import DUMMY_SORT
from fest.symstructs import SymStruct
from fest.symstructs.pickling import PicklableSymStruct
from fest.templates import DomainLike
from fest.templates.model_finder import find_sym_model
from fest.templates.stdlib.arith import TemplateChoice
from fest.z3_utils import find_order_relation
from problems.problem import ProblemCallable


class FlattenedDomain(dict[PicklableSort | None, int]):
    def domain(
        self, pickleable_summary_sort: PicklableSort
    ) -> dict[z3.SortRef, DomainLike]:
        summary_sort = pickleable_summary_sort.sort
        domain: dict[z3.SortRef, DomainLike] = {
            sort.sort: size
            for sort, size in self.items()
            if sort is not None and not sort == pickleable_summary_sort
        }
        domain[summary_sort] = (
            self[pickleable_summary_sort],
            self[None],
        )
        return domain


def find_sym_struct(
    formula: z3.BoolRef,
    template: TemplateChoice,
    domain: FlattenedDomain,
    summary_sort: z3.SortRef | PicklableSort,
) -> SymStruct | None:
    if not isinstance(summary_sort, PicklableSort):
        summary_sort = PicklableSort.create(summary_sort)

    return find_sym_model(
        formula,
        template.get_template().for_domain(domain.domain(summary_sort)),
    )


@dataclass(eq=True, frozen=True, slots=True)
class Task:
    pickleable_formula: PicklableExpr
    template: TemplateChoice
    pickleable_summary_sort: PicklableSort
    flattened_domain: FlattenedDomain

    @property
    def formula(self) -> z3.BoolRef:
        return self.pickleable_formula.formula


type Result = PicklableSymStruct


def _model_finder_runner(task: Task) -> JobResult[Result]:
    try:
        result = find_sym_struct(
            task.formula,
            task.template,
            task.flattened_domain,
            task.pickleable_summary_sort,
        )
        if result is None:
            return None
        return PicklableSymStruct.from_struct(result)
    except Exception as e:
        print(repr(e))
        return None


@dataclass(eq=True, frozen=True)
class Heuristic:
    formula: z3.BoolRef
    template: TemplateChoice
    max_size: int

    @cached_property
    def pickleable_formula(self) -> PicklableExpr:
        return PicklableExpr.create(self.formula)

    @cached_property
    def order(self) -> Relation | None:
        return find_order_relation(self.formula)

    @cached_property
    def formula_sorts(self) -> set[z3.SortRef]:
        return uninterpreted_sorts(self.formula)

    @cached_property
    def pure(self) -> bool:
        return not self.formula_sorts

    @cached_property
    def sorts(self) -> set[z3.SortRef]:
        if self.pure:
            return {z3.DeclareSort(DUMMY_SORT)}
        else:
            return self.formula_sorts

    @cached_property
    def summary_sort(self) -> z3.SortRef:
        if self.order is None:
            sorts = {str(sort): sort for sort in self.sorts}
            for sort in config.fest.heuristic.preferred_summary_sorts(list[str]):
                if sort in sorts:
                    return sorts[sort]

            assert self.sorts
            return next(iter(self.sorts))
        else:
            return self.order.domain(0)

    @cached_property
    def pickleable_summary_sort(self) -> PicklableSort:
        return PicklableSort.create(self.summary_sort)

    @cached_property
    def other_sorts(self) -> set[z3.SortRef]:
        return self.sorts - {self.summary_sort}

    def combinations(self) -> list[FlattenedDomain]:
        if self.pure:
            return [
                FlattenedDomain({PicklableSort.create(self.summary_sort): 1, None: 0})
            ]

        sizes = list(range(1, 1 + self.max_size))

        summary_sizes = list(
            range(config.fest.heuristic.min_summary_size(int), 1 + self.max_size)
        )

        sorts = {PicklableSort.create(sort) for sort in self.sorts}

        sort_options: list[list[tuple[PicklableSort | None, int]]] = [
            [(sort, size) for size in sizes] for sort in sorts
        ]
        sort_options.append([(None, size) for size in summary_sizes])

        return list(map(FlattenedDomain, product(*sort_options)))

    def weight(self, domain: FlattenedDomain) -> int:
        other_domains = [
            size
            for sort, size in domain.items()
            if sort is not None and not sort == self.pickleable_summary_sort
        ]

        weight = sum(other_domains)

        regular = domain[self.pickleable_summary_sort]
        summary = domain[None]

        weight += 4 * regular + 10 * summary

        if self.other_sorts:
            min_size = min(regular, min(other_domains))
            max_size = max(regular, max(other_domains))
            diff = max_size - min_size

            if diff > 1:
                weight += 20

        return weight

    def run(self, *, dry: bool = False, quiet: bool = False) -> SymStruct | None:
        domain_options = sorted(self.combinations(), key=self.weight)

        tasks = []

        for i, domain in enumerate(domain_options):
            if not quiet:
                print(f"Template #{i + 1}")
                print(f"\tDomain is {domain}")
                print(f"\tWeight is {self.weight(domain)}")
                print("---")
            if dry:
                continue

            if config.fest.legacy.disable_parallelism:
                if (
                    model := find_sym_struct(
                        self.formula, self.template, domain, self.summary_sort
                    )
                ) is not None:
                    return model
            else:
                tasks.append(
                    Task(
                        self.pickleable_formula,
                        self.template,
                        self.pickleable_summary_sort,
                        domain,
                    )
                )

        if not config.fest.legacy.disable_parallelism:
            with Pool[Result](
                config.fest.parallel.processes(int), start_upfront=True
            ) as pool:

                for task in tasks:
                    pool.submit(_model_finder_runner, task)
                pool.done()

                try:
                    pickled_result = pool.get(timeout=config.fest.timeout(float))
                    if pickled_result is None:
                        return None

                    return pickled_result.struct
                except queue.Empty:
                    return None

        else:
            return None


@config(
    max_size=config.fest.heuristic.max_size,
)
def find_sym_model_by_heuristic(
    formula: ProblemCallable | z3.BoolRef,
    *,
    dry: bool = False,
    quiet: bool = True,
    template: TemplateChoice = TemplateChoice.STANDARD,
    max_size: int = 1,
) -> SymStruct | None:
    if callable(formula):
        formula = formula()
    return Heuristic(formula, template, max_size).run(dry=dry, quiet=quiet)
