from pprint import pprint
from typing import Any, TypeVar, Dict, List

import act4e_interfaces as I
from .sets_representation import SolFiniteSetRepresentation
from .sets_product import MyFiniteSetProduct

A = TypeVar("A")
B = TypeVar("B")

class MyFiniteMap(I.FiniteMap[A, B]):
    _source: I.FiniteSet[A]
    _target: I.FiniteSet[B]
    _values: List[List]

    def __init__(self, source: I.FiniteSet[A], target: I.FiniteSet[B], values: Dict[A, B]):
        self._source = source
        self._target = target
        self._values = values

    def source(self) -> I.FiniteSet[A]:
        return self._source

    def target(self) -> I.FiniteSet[B]:
        return self._target

    def __call__(self, a: A) -> B:
        # print("Function call")
        # print("src: ", self._source.elements())
        # print("dst: ", self._target.elements())
        # print("val: ", self._values)
        # print("inp: ", a)
        assert a in self._source.elements()

        rv = None
        for v in self._values:
            if self._source.equal(v[0], a):
                # print(f"{v[0]} equals {a}")
                rv = v[1]

        assert self._target.contains(rv)
        # print("ret: ", rv)
        # print("the type of rv is", type(rv))
        # print()

        return rv


class SolFiniteMapRepresentation(I.FiniteMapRepresentation):
    def load(self, h: I.IOHelper, s: I.FiniteMap_desc) -> I.FiniteMap[A, B]:
        fsr = SolFiniteSetRepresentation()

        if not "values" in s.keys():
            raise I.InvalidFormat()

        if len(s["values"]) == 0:
            raise I.InvalidFormat()

        A = fsr.load(h, s["source"])
        B = fsr.load(h, s["target"])
        values = []

        for v in s["values"]:
            a, b = A.load(h, v[0]), B.load(h, v[1])

            if not A.contains(a):
                print(f"{v[0]} is not in A: {A.elements()}")
                raise I.InvalidFormat()

            if not B.contains(b):
                print(f"{v[1]} is not in B: {B.elements()}")
                raise I.InvalidFormat()

            values.append([a, b])

        # Check that the domain values are unique
        domain = []
        for a, _ in values:
            if a in domain:
                raise I.InvalidFormat()
            domain.append(a)

        return MyFiniteMap(A, B, values)


    def save(self, h: I.IOHelper, m: I.FiniteMap[Any, Any]) -> I.FiniteMap_desc:
        fsr = SolFiniteSetRepresentation()
        d = {
            "source": fsr.save(h, m.source()),
            "target": fsr.save(h, m.target()),
            "values": []
        }

        for a in m.source().elements():
            b = m(a)
            d["values"].append([
                m.source().save(h, a),
                m.target().save(h, b)])

        return d