from typing import Any, Dict, List, Optional
import dataclasses
import glob
import os
import sys
import json

import spacy
from spacy.language import Language

from sftp import SpanPredictor


@dataclasses.dataclass
class FrameAnnotation:
    tokens: List[str] = dataclasses.field(default_factory=list)
    pos: List[str] = dataclasses.field(default_factory=list)


@dataclasses.dataclass
class MultiLabelAnnotation(FrameAnnotation):
    frame_list: List[List[str]] = dataclasses.field(default_factory=list)
    lu_list: List[Optional[str]] = dataclasses.field(default_factory=list)

    def to_txt(self):
        for i, tok in enumerate(self.tokens):
            yield f"{tok} {self.pos[i]} {'|'.join(self.frame_list[i]) or '_'} {self.lu_list[i] or '_'}"


def convert_to_seq_labels(sentence: List[str], structures: Dict[int, Dict[str, Any]]) -> List[List[str]]:
    labels = [[] for _ in sentence]

    for struct_id, struct in structures.items():
        tgt_span = struct["target"]
        frame = struct["frame"]

        for i in range(tgt_span[0], tgt_span[1] + 1):
            labels[i].append(f"T:{frame}@{struct_id:02}")
        for role in struct["roles"]:
            role_span = role["boundary"]
            role_label = role["label"]
            for i in range(role_span[0], role_span[1] + 1):
                prefix = "B" if i == role_span[0] else "I"
                labels[i].append(f"{prefix}:{frame}:{role_label}@{struct_id:02}")
    return labels


def predict_combined(
    spacy_model: Language,
    sentences: List[str],
    tgt_predictor: SpanPredictor,
    frm_predictor: SpanPredictor,
    bnd_predictor: SpanPredictor,
    arg_predictor: SpanPredictor,
) -> List[MultiLabelAnnotation]:
    
    annotations_out = []

    for sent_idx, sent in enumerate(sentences):

        sent = sent.strip()

        print(f"Processing sent with idx={sent_idx}: {sent}")
        
        doc = spacy_model(sent)
        sent_tokens = [t.text for t in doc]

        tgt_spans, _, _ = tgt_predictor.force_decode(sent_tokens)

        frame_structures = {}

        for i, span in enumerate(tgt_spans):
            span = tuple(span)
            _, fr_labels, _ = frm_predictor.force_decode(sent_tokens, child_spans=[span])
            frame = fr_labels[0]
            if frame == "@@VIRTUAL_ROOT@@@":
                continue

            boundaries, _, _ = bnd_predictor.force_decode(sent_tokens, parent_span=span, parent_label=frame)
            _, arg_labels, _ = arg_predictor.force_decode(sent_tokens, parent_span=span, parent_label=frame, child_spans=boundaries)

            frame_structures[i] = {
                "target": span,
                "frame": frame,
                "roles": [
                    {"boundary": bnd, "label": label}
                    for bnd, label in zip(boundaries, arg_labels)
                    if label != "Target"
                ]
            }
        annotations_out.append(MultiLabelAnnotation(
            tokens=sent_tokens,
            pos=[t.pos_ for t in doc],
            frame_list=convert_to_seq_labels(sent_tokens, frame_structures),
            lu_list=[None for _ in sent_tokens]
        ))
    return annotations_out


def main(input_folder):

    print("Loading spaCy model ...")
    nlp = spacy.load("it_core_news_md")

    print("Loading predictors ...")
    zs_predictor = SpanPredictor.from_path("/data/p289731/cloned/lome-models/models/spanfinder/model.mod.tar.gz", cuda_device=0)
    ev_predictor = SpanPredictor.from_path("/scratch/p289731/lome-training-files/train-evalita-plus-fn-vanilla/model.tar.gz", cuda_device=0)


    print("Reading input files ...")
    for file in glob.glob(os.path.join(input_folder, "*.txt")):
        print(file)
        with open(file, encoding="utf-8") as f:
            sentences = list(f)

        annotations = predict_combined(nlp, sentences, zs_predictor, ev_predictor, ev_predictor, ev_predictor)

        out_name = os.path.splitext(os.path.basename(file))[0]
        with open(f"../../data-out/{out_name}.combined_zs_ev.tc_bilstm.txt", "w", encoding="utf-8") as f_out:
            for ann in annotations:
                for line in ann.to_txt():
                    f_out.write(line + os.linesep)
                f_out.write(os.linesep)

        with open(f"../../data-out/{out_name}.combined_zs_ev.tc_bilstm.json", "w", encoding="utf-8") as f_out:
            json.dump([dataclasses.asdict(ann) for ann in annotations], f_out)


if __name__ == "__main__":
    main(sys.argv[1])