from sage.matroids.advanced import *

#Potentially faster than sage.matroids.advanced.get_nonisomorphic_matroids()
def _get_nonisomorphic_matroids(matroids):
    return {MatroidIsomorphismClass(m) : m for m in matroids}.values()

class MatroidIsomorphismClass:
    def __init__(self, matroid):
        if isinstance(matroid, BasisMatroid):
            self.matroid = matroid
        else:
            self.matroid = BasisMatroid(matroid)

    def __eq__(self, other):
        return self.matroid.is_isomorphic(other.matroid)

    def __ne__(self, other):
        return not self.matroid.is_isomorphic(other.matroid)
        
    def __hash__(self):
        return hash(self.matroid._bases_invariant())

#a catalog of matroids, stored on disk, keyed by isomorphism class.
class MICCatalog:
    def __init__(self, directory, bug_workaround=True):
        self.directory = directory
        self.bug_workaround = bug_workaround

    #load the Mayhew-Royle 9-element matroid catalog, specified as a file
    def setup_MR_catalog(self, filename):
        with open(filename) as f:
            curr_matroids = {}
            curr_rank = 0
            curr_corank = 0
            for line in f.readlines():
                [number,content]=line.strip().split('\t')
                matroid = self._convert_MR_string_to_matroid(content)

                #save as we go (so don't need to keep too many in memory)
                if matroid.full_rank() != curr_rank or matroid.full_corank() != curr_corank:
                    save(curr_matroids, self.directory + "/" + self._get_filename(curr_rank, curr_corank))
                    curr_matroids = {}
                    curr_rank = matroid.full_rank()
                    curr_corank = matroid.full_corank()

                matroidId = int(number)
                curr_matroids[matroidId] = matroid
                
            #final save
            save(curr_matroids, self.directory + "/" + self._get_filename(curr_rank, curr_corank))

    #saving MICs is slow -- it is quicker to rebuild them.
    def _load(self, filename):
        #we call .bases() on each matroid, as a bug workaround
        # (I believe the bug has been corrected in a recent version)
        pairs_iter = load(filename).iteritems()
        return {MatroidIsomorphismClass(matroid) : key for (key, matroid) in pairs_iter if matroid.bases() != None}

    def _get_filename(self, rk, cork):
        n = rk + cork
        return "mic_dicts" +str(n)+ "_" + str(rk) + ".sobj"

    #return the value stored for the given matroid isomorphism class
    # for the Mayhew-Royle 9-element matroid catalog, this is the ID.
    def lookup(self, matroid):
        rank = matroid.full_rank()
        corank = matroid.full_corank()
        dicts = self._load_mic_dicts_rc(rank, corank)
        return dicts[MatroidIsomorphismClass(matroid)]

    def _from_id(self, idnumber, filename):
        with open(filename) as f:
            for line in f.readlines():
                [number,content]=line.strip().split('\t')
                if int(number) == idnumber:
                    return self._convert_MR_string_to_matroid(content)
        return None

    #return the matroid with the given rank, corank, and "value"
    #(for the Mayhew-Royle 9-element matroid catalog, the "value" is the ID).
    def get(self, rk, corank, idnumber):
        for (k, v) in self._load_mic_dicts_rc(rk, corank).items():
            if v == idnumber:
                return k.matroid

    #return the matroids of the given rank and corank, having one of the given "value"s
    #(for the Mayhew-Royle 9-element matroid catalog, the "value" is the ID).
    def get_all(self, rk, corank, idnumbers):
        idnumbers_set = set(idnumbers)
        result_dict = {}
        for (k, v) in self._load_mic_dicts_rc(rk, corank).items():
            if v in idnumbers_set:
                result_dict[v] = k.matroid
        
        return [result_dict.get(idnum) for idnum in idnumbers]

    def lookup_list(self, matroids):
        result_map = {}
        keys = set()
        for mat in matroids:
            keys.add((mat.full_rank(), mat.full_corank()))

        for (r,c) in keys:
            dicts = self._load_mic_dicts_rc(r, c)
            for i,mat in enumerate(matroids):
                if mat.full_rank() == r and mat.full_corank() == c:
                    result_map[i] = dicts[MatroidIsomorphismClass(mat)]
        return [result_map[i] for i in sorted(result_map)]

    def _convert_MR_string_to_matroid(self, content):
        v = content.split(' ')
        n = int(v[0])
        r = int(v[1])
        nb = int(v[2])
        bases_list = v[3:]

        e_string = ""
        for i in range(n):
            e_string = e_string + str(i)

        bs_strings = []
        for bs in bases_list:
            bs = bs.strip("'")
            bs_strings.append(bs)
        m = Matroid(groundset=e_string, bases=bs_strings)
        assert m.full_rank() == r and len(m.bases()) == nb
        return m

    def _load_mic_dicts_rc(self, rk, cork):
        filename = self._get_filename(rk, cork)
        if filename in os.listdir(self.directory):
            matroids_rc = self._load(self.directory + "/" + filename)
        else:
            matroids_rc = {}
        
        if self.bug_workaround:
            for mic in matroids_rc.keys():
                _ = mic.matroid.bases()

        return matroids_rc