#python merge.py parent.spm child.spm merged.spm
import sys
import sentencepiece as spm
import yaml
import sentencepiece_model_pb2 as model
def merge_models():
    m_parent = model.ModelProto()
    m_parent.ParseFromString(open('/home/large/data/models/marian/daen/corp/vocab.daen.spm', 'rb').read())
    m_child = model.ModelProto()
    m_child.ParseFromString(open('/home/large/data/models/marian/enno.20200501/corp/vocab.enno.spm', 'rb').read())
    m_new = model.ModelProto()
    m_new.ParseFromString(open('/home/large/data/models/marian/enno.20200501/corp/vocab.enno.spm', 'rb').read())
    parent_piece_to_id={}
    for i,p in enumerate(m_parent.pieces):
        parent_piece_to_id[p.piece]=i

    used_ids=[]
    unique_child_pieces=[]
    for i,p in enumerate(m_child.pieces):
        if p.piece in parent_piece_to_id:
            m_new.pieces[parent_piece_to_id[p.piece]].piece=p.piece
            used_ids.append(parent_piece_to_id[p.piece])
        else:
            unique_child_pieces.append(p)
    i=0
    for pos in range(len(m_new.pieces)):
        if not pos in used_ids:
            m_new.pieces[pos].piece=unique_child_pieces[i].piece
            i+=1
    with open('/home/large/data/models/marian/enno.20200501/corp/vocab.enno_matched_daen.spm', 'wb') as f:
        f.write(m_new.SerializeToString())


# takes vocabs generated by spm_export_vocab, merge them and save them in marian .yml format
def merge_vocabs():
    with open("/home/big_maggie/usr/nmt_scripts/spm_match/sentencepiece/src/vocab.enno") as child, \
        open("/home/big_maggie/usr/nmt_scripts/spm_match/sentencepiece/src/vocab.daen") as parent, \
        open("/home/big_maggie/usr/nmt_scripts/spm_match/sentencepiece/src/vocab.enno_daen_merged_frequency.yml","w") as merged:
        parent_vocab=[line.split('\t')[0] for line in parent]
        child_vocab=[line.split('\t')[0] for line in child]
        parent_subwords=set(parent_vocab)
        child_subwords=set(child_vocab)
        common_subwords=parent_subwords & child_subwords
        unique_child_subwords=[s for s in child_vocab if s not in parent_vocab]
        merged_vocab=[None]*len(child_vocab)
        #print(common_subwords)
        print(len(common_subwords))
        print(len(unique_child_subwords))
        used_ids=[]
        for subword in common_subwords:
            merged_vocab[parent_vocab.index(subword)]=subword
            used_ids.append(child_vocab.index(subword))
        i=0
        for position,subword in enumerate(merged_vocab):
            if subword==None:
                merged_vocab[position]=unique_child_subwords[i]
                i+=1
        print(i)
        print(merged_vocab==child_vocab)
        for i,s in enumerate(merged_vocab):
            s=s.replace('"','\\"')
            if not s.isalpha() and not s=="</s>" and not s=="<unk>":
                s='"{}"'.format(s)
            merged.write("{}: {}\n".format(s,i))
    #print(merged_subwords)

# takes vocabs generated by marian (parent) and spm model (child), merge them and save them in marian .yml format
def merge_vocabs_yml_and_spm():
    with open("/home/large/data/models/marian/encz_exp/wmt19/corp/vocab.encs.yml",errors="ignore") as parent:
        m_child = model.ModelProto()
        m_child.ParseFromString(open('/home/large/data/models/marian/enlt.20200524/corp/vocab.lten.spm', 'rb').read())
        m_new = model.ModelProto()
        m_new.ParseFromString(open('/home/large/data/models/marian/enlt.20200524/corp/vocab.lten.spm', 'rb').read())
        parent_piece_to_id = {}
        parent_piece_to_id = yaml.load(parent, Loader=yaml.FullLoader)
        used_ids = []
        unique_child_pieces = []
        for i, p in enumerate(m_child.pieces):
            if p.piece in parent_piece_to_id:
                m_new.pieces[parent_piece_to_id[p.piece]].piece = p.piece
                used_ids.append(parent_piece_to_id[p.piece])
            else:
                unique_child_pieces.append(p)
        i = 0
        for pos in range(len(m_new.pieces)):
            if not pos in used_ids:
                m_new.pieces[pos].piece = unique_child_pieces[i].piece
                i += 1
        print(len(unique_child_pieces))
        print(i)
        print(m_new==m_child)

        with open('/home/large/data/models/marian/enlt.20200524/corp/vocab.lten.czen.merged.spm', 'wb') as f:
            f.write(m_new.SerializeToString())


merge_vocabs_yml_and_spm()
