from core.mind.rl_mind import AlgorithmicMind, MindOutput
from core.mind.cognitive_actions import (
    CognitiveAction, ActionParams, ActionResult, ActionHandler,
)
from core.mind.mind_state import MindState
from core.mind.solving_loop import MindSolvingLoop, ThinkingSession, UniversalSolution
from core.mind.knowledge_graph import AlgorithmicKnowledgeGraph
from core.mind.safety_manager import MindSafetyManager
from core.mind.model_config import MindConfig, DEFAULT_CONFIG


def create_mind(
    config: MindConfig | None = None,
    solver=None,
    synthesizer=None,
    persist_dir: str = "~/.aalgoi",
) -> MindSolvingLoop:
    from pathlib import Path
    from core.mind.adapters import (
        ExecutorAdapter, SynthesizerAdapter, ProverAdapter,
    )
    from core.reasoning.comprehension_engine import DeepComprehensionEngine
    from core.reasoning.correctness_prover import CorrectnessProver

    cfg = config or DEFAULT_CONFIG

    mind = AlgorithmicMind(cfg)

    kg_path = Path(persist_dir).expanduser() / "kg"
    kg = AlgorithmicKnowledgeGraph(kg_path)

    executor = ExecutorAdapter(solver)
    synth = SynthesizerAdapter(synthesizer)
    prover = ProverAdapter(CorrectnessProver())

    comprehension = DeepComprehensionEngine()

    handler = ActionHandler(
        kg=kg,
        executor=executor,
        synthesizer=synth,
        prover=prover,
        comprehension=comprehension,
    )

    safety = MindSafetyManager(Path(persist_dir).expanduser())

    weights_path = Path(persist_dir).expanduser() / "personal_model.pt"
    if weights_path.exists():
        import torch
        mind.load_state_dict(torch.load(weights_path, map_location="cpu"))
    else:
        base_path = Path(__file__).parent / "pretrained" / "base_v2.pt"
        if base_path.exists():
            import torch
            mind.load_state_dict(torch.load(base_path, map_location="cpu"))

    loop = MindSolvingLoop(
        mind=mind,
        kg=kg,
        action_handler=handler,
        safety=safety,
    )

    return loop


__all__ = [
    "AlgorithmicMind", "MindOutput",
    "CognitiveAction", "ActionParams", "ActionResult", "ActionHandler",
    "MindState",
    "MindSolvingLoop", "ThinkingSession", "UniversalSolution",
    "AlgorithmicKnowledgeGraph",
    "MindSafetyManager",
    "MindConfig", "DEFAULT_CONFIG",
    "create_mind",
]
