
import json
from collections import defaultdict

import tarski
from tarski.io import PDDLReader
from tarski.io import fstrips as iofs
from tarski.syntax.formulas import is_atom
from tarski.syntax.transform.action_grounding import ground_schema_into_plain_operator_from_grounding
from tarski.grounding.common import StateVariableLite
from tarski.grounding.lp_grounding import LPGroundingStrategy
from tarski.util import SymbolIndex

from pddl.core import Problem
from pddl.parser.domain import DomainParser
from pddl.parser.problem import ProblemParser

def entails(state, partialstate):
    return partialstate <= state

def progress(state, act):
    assert entails(state, act.pres), "Cannot progress with inconsistent state / action precondition:\n\t Action: "+act.name+"\n\t State: \n\t\t"+'\n\t\t'.join(state)
    return (state - act.dels) | act.adds

def regress(state, act):
    assert len(state & act.dels) == 0, f"Cannot regress with inconsistent state / action delete effect:\n\t Action: "+act.name+"\n\t State: \n\t\t"+'\n\t\t'.join(state)
    return (state - act.adds) | act.pres



def fix_name(s):
    # (act param)
    if '(' == s[0] and ')' == s[-1]:
        return s[1:-1]
    # make it space separated
    s = s.replace(', ', ' ').replace(',', ' ')
    # act(param)
    if '(' in s:
        assert ')' == s[-1], f"Broken name? {s}"
        s = s.replace('(', ' ').replace(')', '')
    # act param
    return s


class Action:
    def __init__(self, name, pre, add, delete):
        self.name = name
        self.pres = pre
        self.adds = add
        self.dels = delete

    def __str__(self):
        pres = "{" + ", ".join([f"({a})" for a in self.pres]) + "}"
        adds = "{" + ", ".join([f"({a})" for a in self.adds]) + "}"
        dels = "{" + ", ".join([f"({a})" for a in self.dels]) + "}"

        return f"< {self.name}, {pres}, {adds}, {dels} >"

    def toJSON(self):
        return json.dumps(
            { "name": self.name, "preconditions": [f"({a})" for a in self.pres], "add_effects": [f"({a})" for a in self.adds], "delete_effects": [f"({a})" for a in self.dels] },
            sort_keys=True,
            indent=4)

    def __repr__(self):
        return self.name
    
    def __eq__(self, action):
        return self.name == action.name
    

    def __hash__(self):
        return hash(self.name)

class STRIPS:
    def __init__(self, domain, problem):
        self.domain_file = domain 
        self.problem_file = problem
        self.reader = PDDLReader(raise_on_error=True)
        self.reader.parse_domain(domain)
        self.problem = self.reader.parse_instance(problem)
        (self.grounded_fluents, init, goal, self.operators, self.grounder) = self.ground_problem(self.problem)

        self.fluents = set([fix_name(str(f)) for f in self.grounded_fluents])
        self.fluents_map = dict()
        for f in self.grounded_fluents:
            self.fluents_map[fix_name(str(f))] = f
        self.init = set([fix_name(str(f)) for f in init])
        self.goal = set([fix_name(str(f)) for f in goal])
        self.actions = set()
        self.action_map = {}
        self.init_fluents = [self.fluents_map[f] for f in self.init]

        self.static_predicates = [ i.name for i in self.grounder.static_symbols]
        for op in self.operators:
            act = self.operator_to_action(op)
            self.actions.add(act)
            self.action_map[act.name.lower()] = act


    def __str__(self):
        fluents = "P = {" + ", ".join([f"({a})" for a in self.fluents]) + "}"
        init = "I = {" + ", ".join([f"({a})" for a in self.init]) + "}"
        goal = "G = {" + ", ".join([f"({a})" for a in self.goal]) + "}"
        actions = "A = {" + "\n ".join([a.__str__() for a in self.actions]) + "}"
        return fluents + ",\n" + init + "\n" + goal + "\n" + actions

    def toJSON(self):
        actions = [a.toJSON() for a in self.actions]
        return json.dumps( 
            { "fluents": list(self.fluents), "initial_state": list(self.init), "goal": list(self.goal), "actions": actions },
            sort_keys=True,
            indent=4)

    def operator_to_action(self, op, check_fluents= True, check_static=False):
        adds = {fix_name(str(f.atom)) for f in op.effects if isinstance(f, iofs.AddEffect)} & self.fluents
        dels = {fix_name(str(f.atom)) for f in op.effects if isinstance(f, iofs.DelEffect)} & self.fluents
        pre = self.fix_pre_name(op.precondition) 
        if check_fluents:
            pre = pre & self.fluents
        if check_static:
            pre = {p for p in pre if p.split()[0] not in self.static_predicates}
        act = Action(fix_name(str(op)), pre, adds, dels)
        return act

    def fix_pre_name(self, precondition):
        if not is_atom(precondition):
            return {fix_name(str(f)) for f in precondition.subformulas}
        return {fix_name(str(precondition))} 

    def action(self, name):
        return self.action_map[fix_name(name).lower()]

    def get_action_or_none(self, name):
        if '(' in name and ')' != name[-1]:
            return None
        return self.action_map.get(fix_name(name).lower(), None)

    def fluent(self, name):
        return fix_name(name)
    
    def static_symbols(self):
        return list(self.grounder.static_symbols)
    
    def fluent_symbols(self):
        return list(self.grounder.fluent_symbols)


    def get_grounded_atoms(self, symbol):
        variables = SymbolIndex()
        lang = symbol.language
        key = 'atom_' + symbol.name
        model = self.grounder._solve_lp()
        if key in model:  # in case there is no reachable ground state variable from that fluent symbol
            for binding in model[key]:
                binding_with_constants = tuple(lang.get(c) for c in binding)
                variables.add(StateVariableLite(symbol, binding_with_constants))
        return variables

    def get_applicable_actions(self, s):
        return [a for a in self.actions if entails(s, a.pres)]

    
    def ground_problem(self, problem):
        grounder = LPGroundingStrategy(problem, include_variable_inequalities=True)
        action_groundings = grounder.ground_actions()
        operators = []
        for action_name, groundings in action_groundings.items():
            action = problem.get_action(action_name)
            for grounding in groundings:
                # print(type(grounding[0]), grounding)
                operators.append(ground_schema_into_plain_operator_from_grounding(action, grounding))

        grounded_fluents = set([grounded_fluent.to_atom() for grounded_fluent in grounder.ground_state_variables().objects])
        init = [f for f in problem.init.as_atoms() if f in grounded_fluents]
        if isinstance(problem.goal, tarski.syntax.Atom):
            goal = [problem.goal]
        else:
            goal = [f for f in problem.goal.subformulas if f in grounded_fluents]

        return (grounded_fluents, init, goal, operators, grounder)


    def get_static(self):
        static_symbols = self.static_symbols()
        ret = []
        for symbol in static_symbols:
            ret.extend(self.get_grounded_atoms(symbol))
        return set([fix_name(str(x)) for x in ret])

    def PDDL_replace_init_pddl_parser(self, s):
        d = DomainParser()(open(self.domain_file, "r").read().lower())
        p = ProblemParser()(open(self.problem_file, "r").read().lower())

        new_state = get_atoms_pddl(d, p, s | self.get_static())

        new_p = Problem(p.name, domain=d,
                    objects=p.objects,
                    init=new_state,
                    goal=p.goal)

        return d, new_p


def get_atoms_pddl(d, p, atoms):
    # print(atoms)
    objs = set()
    preds = defaultdict(list)
    for atom in atoms:
        a = atom.lower().strip().split(" ")
        args = a[1:]
        # print(f"atom name |{a[0]}|")
        preds[a[0]].append(args)
        objs |= set(args)
    # print(preds)

    constants = [o for o in p.objects | d.constants if o.name.lower() in objs]
    constants_dict = {}
    for c in constants:
        constants_dict[c.name.lower()] = c
    assert len(objs) == len(constants), f"Could not identify all objects: {objs - set(constants_dict.keys())} not found, {set(constants_dict.keys()) - objs} should not be there"

    state = []
    covered_preds = set()
    for f in d.predicates:
        name = f.name.lower()
        # print(f"Checking predicate |{name}|")
        if name in preds:
            covered_preds.add(name)
            assert len(preds[name][0]) == f.arity, f"The arity does not match: {preds[name]} vs {f.terms}"
            # Going over the lists of objects, adding ground predicate for each 
            # print(f.name, preds[f.name])
            for ob in preds[name]:
                c = [ constants_dict[o] for o in ob]
                # print("f: ", f, type(f), [x.type_tags for x in f.terms])
                # print("c: ", c, [type(x) for x in c], [x.type_tags for x in c])
                state.append(f(*c))

    assert len(covered_preds) == len(preds.keys()), f"Covered predicates: \n{sorted(list(covered_preds))} vs \n{sorted(list(preds.keys()))}"        
    return set(state)
