import torch
import pickle
import gzip

from torchtext import data
from slt.signjoey.model import build_model
from slt.signjoey.batch import Batch
from slt.signjoey.data import make_data_iter
from slt.signjoey.vocabulary import PAD_TOKEN
from slt.signjoey.dataset import SignTranslationDataset
from slt.signjoey.phoenix_utils.phoenix_cleanup import (
    clean_phoenix_2014,
    clean_phoenix_2014_trans,
)
from slt.signjoey.vocabulary import (
    build_vocab,
    UNK_TOKEN,
    EOS_TOKEN,
    BOS_TOKEN,
    PAD_TOKEN,
)
from slt.signjoey.helpers import (
    bpe_postprocess,
    load_config,
    load_checkpoint,
)

torch.backends.cudnn.deterministic = True


def translate(cfg_file: str, ckpt: str, input_path: str) -> str:
    cfg = load_config(cfg_file)

    batch_type = cfg["training"].get("batch_type", "sentence")
    use_cuda = cfg["training"].get("use_cuda", False)
    level = cfg["data"]["level"]
    dataset_version = cfg["data"].get("version", "phoenix_2014_trans")
    translation_max_output_length = cfg["training"].get(
        "translation_max_output_length", None
    )
    txt_lowercase = cfg["data"]["txt_lowercase"]
    
    best_recog_beam_size = cfg["translating"]["best_recog_beam_size"]
    best_trans_beam_size = cfg["translating"]["best_trans_beam_size"]
    best_trans_beam_alpha = cfg["translating"]["best_trans_beam_alpha"]
    sgn_dim = cfg["data"]["feature_size"]

    do_recognition = cfg["training"].get("recognition_loss_weight", 1.0) > 0.0
    do_translation = cfg["training"].get("translation_loss_weight", 1.0) > 0.0

    gls_vocab = build_vocab(
        field="gls",
        min_freq=None,
        max_size=None,
        dataset=None,
        vocab_file='slt/data/gls.vocab',
    )

    txt_vocab = build_vocab(
        field="txt",
        min_freq=None,
        max_size=None,
        dataset=None,
        vocab_file='slt/data/txt.vocab',
    )

    def tokenize_text(text):
        return list(text) if level == "char" else text.split()

    def tokenize_features(features):
        ft_list = torch.split(features, 1, dim=0)
        return [ft.squeeze() for ft in ft_list]

    def stack_features(features, something):
        return torch.stack([torch.stack(ft, dim=0) for ft in features], dim=0)

    sequence_field = data.RawField()
    signer_field = data.RawField()

    sgn_field = data.Field(
        use_vocab=False,
        init_token=None,
        dtype=torch.float32,
        preprocessing=tokenize_features,
        tokenize=lambda features: features,
        batch_first=True,
        include_lengths=True,
        postprocessing=stack_features,
        pad_token=torch.zeros((sgn_dim,)),
    )

    gls_field = data.Field(
        pad_token=PAD_TOKEN,
        tokenize=tokenize_text,
        batch_first=True,
        lower=False,
        include_lengths=True,
    )

    txt_field = data.Field(
        init_token=BOS_TOKEN,
        eos_token=EOS_TOKEN,
        pad_token=PAD_TOKEN,
        tokenize=tokenize_text,
        unk_token=UNK_TOKEN,
        batch_first=True,
        lower=txt_lowercase,
        include_lengths=True,
    )

    gls_field.vocab = gls_vocab
    txt_field.vocab = txt_vocab

    trans_path = "slt/data/trans.pkl"
    build_embeddings(trans_path, input_path)

    trans_data = SignTranslationDataset(
        path=trans_path,
        fields=(sequence_field, signer_field, sgn_field, gls_field, txt_field),
    )
    
    trans_iter = make_data_iter(
        dataset=trans_data,
        batch_size=1,
        batch_type=batch_type,
        shuffle=False,
        train=False,
    )

    model = build_model(
        cfg=cfg["model"],
        gls_vocab=gls_vocab,
        txt_vocab=txt_vocab,
        sgn_dim=sgn_dim,
        do_recognition=False,
        do_translation=True,
    )
    
    ckpt = cfg["training"]["load_model"]
    model_checkpoint = load_checkpoint(ckpt, use_cuda=use_cuda)
    model.load_state_dict(model_checkpoint["model_state"])

    if use_cuda: model.cuda()
    model.eval()

    with torch.no_grad():
        all_gls_outputs = []
        all_txt_outputs = []
        batch = Batch(
            is_train=False,
            torch_batch=next(iter(trans_iter)),
            txt_pad_index=txt_vocab.stoi[PAD_TOKEN],
            sgn_dim=sgn_dim,
            use_cuda=use_cuda,
            frame_subsampling_ratio=None,
        ) 
        sort_reverse_index = batch.sort_by_sgn_lengths()

        (
            batch_gls_predictions,
            batch_txt_predictions,
            _,
        ) = model.run_batch(
                batch=batch,
                recognition_beam_size=best_recog_beam_size if do_recognition else None,
                translation_beam_size=best_trans_beam_size if do_translation else None,
                translation_beam_alpha=best_trans_beam_alpha if do_translation else None,
                translation_max_output_length=translation_max_output_length if do_translation else None,
            )

        # sort outputs back to original order
        if do_recognition:
            all_gls_outputs.extend(
                [batch_gls_predictions[sri] for sri in sort_reverse_index]
            )
        if do_translation:
            all_txt_outputs.extend(
                batch_txt_predictions[sort_reverse_index]
            )

        if do_recognition:
            assert len(all_gls_outputs) == len(trans_data)
            # decode back to symbols
            decoded_gls = model.gls_vocab.arrays_to_sentences(arrays=all_gls_outputs)

            # Gloss clean-up function
            if dataset_version == "phoenix_2014_trans":
                gls_cln_fn = clean_phoenix_2014_trans
            elif dataset_version == "phoenix_2014":
                gls_cln_fn = clean_phoenix_2014
            else:
                raise ValueError("Unknown Dataset Version: " + dataset_version)

            # Construct gloss sequences for metrics
            gls_hyp = [gls_cln_fn(" ".join(t)) for t in decoded_gls]

        if do_translation:
            assert len(all_txt_outputs) == len(trans_data)
            # decode back to symbols
            decoded_txt = model.txt_vocab.arrays_to_sentences(arrays=all_txt_outputs)
            # evaluate with metric on full dataset
            join_char = " " if level in ["word", "bpe"] else ""
            # Construct text sequences for metrics
            txt_hyp = [join_char.join(t) for t in decoded_txt]
            # post-process
            if level == "bpe":
                txt_hyp = [bpe_postprocess(v) for v in txt_hyp]

    results = {
        "gls_hyp": gls_hyp if do_recognition else None,
        "txt_hyp": txt_hyp if do_translation else None,
    }

    return results['txt_hyp'][0]


def build_embeddings(trans_path: str, input_path: str) -> None:
    with open(trans_path, "rb") as f:
        trans = pickle.load(f)
    
    with gzip.open("slt/data/phoenix14t.pami0.test", "rb") as f:
        annotations = pickle.load(f)

    video_name = 'test/' + input_path.split(sep='/')[2]
    video_name = video_name[:-4]

    for annot in annotations:
        if video_name == annot['name']: 
            trans[0]['sign'] = annot['sign']

    with open(trans_path, "wb") as f:
        pickle.dump(trans, f)
import torch
import pickle
import gzip

from torchtext import data
from slt.signjoey.model import build_model
from slt.signjoey.batch import Batch
from slt.signjoey.data import make_data_iter
from slt.signjoey.vocabulary import PAD_TOKEN
from slt.signjoey.dataset import SignTranslationDataset
from slt.signjoey.phoenix_utils.phoenix_cleanup import (
    clean_phoenix_2014,
    clean_phoenix_2014_trans,
)
from slt.signjoey.vocabulary import (
    build_vocab,
    UNK_TOKEN,
    EOS_TOKEN,
    BOS_TOKEN,
    PAD_TOKEN,
)
from slt.signjoey.helpers import (
    bpe_postprocess,
    load_config,
    load_checkpoint,
)

torch.backends.cudnn.deterministic = True


def translate(cfg_file: str, ckpt: str, input_path: str) -> str:
    cfg = load_config(cfg_file)

    batch_type = cfg["training"].get("batch_type", "sentence")
    use_cuda = cfg["training"].get("use_cuda", False)
    level = cfg["data"]["level"]
    dataset_version = cfg["data"].get("version", "phoenix_2014_trans")
    translation_max_output_length = cfg["training"].get(
        "translation_max_output_length", None
    )
    txt_lowercase = cfg["data"]["txt_lowercase"]
    
    best_recog_beam_size = cfg["translating"]["best_recog_beam_size"]
    best_trans_beam_size = cfg["translating"]["best_trans_beam_size"]
    best_trans_beam_alpha = cfg["translating"]["best_trans_beam_alpha"]
    sgn_dim = cfg["data"]["feature_size"]

    do_recognition = cfg["training"].get("recognition_loss_weight", 1.0) > 0.0
    do_translation = cfg["training"].get("translation_loss_weight", 1.0) > 0.0

    gls_vocab = build_vocab(
        field="gls",
        min_freq=None,
        max_size=None,
        dataset=None,
        vocab_file='slt/data/gls.vocab',
    )

    txt_vocab = build_vocab(
        field="txt",
        min_freq=None,
        max_size=None,
        dataset=None,
        vocab_file='slt/data/txt.vocab',
    )

    def tokenize_text(text):
        return list(text) if level == "char" else text.split()

    def tokenize_features(features):
        ft_list = torch.split(features, 1, dim=0)
        return [ft.squeeze() for ft in ft_list]

    def stack_features(features, something):
        return torch.stack([torch.stack(ft, dim=0) for ft in features], dim=0)

    sequence_field = data.RawField()
    signer_field = data.RawField()

    sgn_field = data.Field(
        use_vocab=False,
        init_token=None,
        dtype=torch.float32,
        preprocessing=tokenize_features,
        tokenize=lambda features: features,
        batch_first=True,
        include_lengths=True,
        postprocessing=stack_features,
        pad_token=torch.zeros((sgn_dim,)),
    )

    gls_field = data.Field(
        pad_token=PAD_TOKEN,
        tokenize=tokenize_text,
        batch_first=True,
        lower=False,
        include_lengths=True,
    )

    txt_field = data.Field(
        init_token=BOS_TOKEN,
        eos_token=EOS_TOKEN,
        pad_token=PAD_TOKEN,
        tokenize=tokenize_text,
        unk_token=UNK_TOKEN,
        batch_first=True,
        lower=txt_lowercase,
        include_lengths=True,
    )

    gls_field.vocab = gls_vocab
    txt_field.vocab = txt_vocab

    trans_path = "slt/data/trans.pkl"
    build_embeddings(trans_path, input_path)

    trans_data = SignTranslationDataset(
        path=trans_path,
        fields=(sequence_field, signer_field, sgn_field, gls_field, txt_field),
    )
    
    trans_iter = make_data_iter(
        dataset=trans_data,
        batch_size=1,
        batch_type=batch_type,
        shuffle=False,
        train=False,
    )

    model = build_model(
        cfg=cfg["model"],
        gls_vocab=gls_vocab,
        txt_vocab=txt_vocab,
        sgn_dim=sgn_dim,
        do_recognition=False,
        do_translation=True,
    )
    
    ckpt = cfg["training"]["load_model"]
    model_checkpoint = load_checkpoint(ckpt, use_cuda=use_cuda)
    model.load_state_dict(model_checkpoint["model_state"])

    if use_cuda: model.cuda()
    model.eval()

    with torch.no_grad():
        all_gls_outputs = []
        all_txt_outputs = []
        batch = Batch(
            is_train=False,
            torch_batch=next(iter(trans_iter)),
            txt_pad_index=txt_vocab.stoi[PAD_TOKEN],
            sgn_dim=sgn_dim,
            use_cuda=use_cuda,
            frame_subsampling_ratio=None,
        ) 
        sort_reverse_index = batch.sort_by_sgn_lengths()

        (
            batch_gls_predictions,
            batch_txt_predictions,
            _,
        ) = model.run_batch(
                batch=batch,
                recognition_beam_size=best_recog_beam_size if do_recognition else None,
                translation_beam_size=best_trans_beam_size if do_translation else None,
                translation_beam_alpha=best_trans_beam_alpha if do_translation else None,
                translation_max_output_length=translation_max_output_length if do_translation else None,
            )

        # sort outputs back to original order
        if do_recognition:
            all_gls_outputs.extend(
                [batch_gls_predictions[sri] for sri in sort_reverse_index]
            )
        if do_translation:
            all_txt_outputs.extend(
                batch_txt_predictions[sort_reverse_index]
            )

        if do_recognition:
            assert len(all_gls_outputs) == len(trans_data)
            # decode back to symbols
            decoded_gls = model.gls_vocab.arrays_to_sentences(arrays=all_gls_outputs)

            # Gloss clean-up function
            if dataset_version == "phoenix_2014_trans":
                gls_cln_fn = clean_phoenix_2014_trans
            elif dataset_version == "phoenix_2014":
                gls_cln_fn = clean_phoenix_2014
            else:
                raise ValueError("Unknown Dataset Version: " + dataset_version)

            # Construct gloss sequences for metrics
            gls_hyp = [gls_cln_fn(" ".join(t)) for t in decoded_gls]

        if do_translation:
            assert len(all_txt_outputs) == len(trans_data)
            # decode back to symbols
            decoded_txt = model.txt_vocab.arrays_to_sentences(arrays=all_txt_outputs)
            # evaluate with metric on full dataset
            join_char = " " if level in ["word", "bpe"] else ""
            # Construct text sequences for metrics
            txt_hyp = [join_char.join(t) for t in decoded_txt]
            # post-process
            if level == "bpe":
                txt_hyp = [bpe_postprocess(v) for v in txt_hyp]

    results = {
        "gls_hyp": gls_hyp if do_recognition else None,
        "txt_hyp": txt_hyp if do_translation else None,
    }

    return results['txt_hyp'][0]


def build_embeddings(trans_path: str, input_path: str) -> None:
    with open(trans_path, "rb") as f:
        trans = pickle.load(f)
    
    with gzip.open("slt/data/phoenix14t.pami0.test", "rb") as f:
        annotations = pickle.load(f)

    video_name = 'test/' + input_path.split(sep='/')[2]
    video_name = video_name[:-4]

    for annot in annotations:
        if video_name == annot['name']: 
            trans[0]['sign'] = annot['sign']

    with open(trans_path, "wb") as f:
        pickle.dump(trans, f)
# sort outputs back to original order # decode back to symbols # Gloss clean-up function # Construct gloss sequences for metrics # decode back to symbols # evaluate with metric on full dataset # Construct text sequences for metrics # post-process