import Frame from "./Frame";
import Role from './Role';
import SrlJsonFormat from "./SrlJsonFormat";

class Srl {
    frameLabelSetName?: string;
    tokens: string[];
    frames: Frame[];
    constructor(tokens: string[], frames: Frame[], frameLabelSetName?:string) {
        this.tokens = tokens || [];
        this.frames = frames || [];
        this.frameLabelSetName = frameLabelSetName;
        Object.freeze(this);
    }

    static parseXml(root:Element):Srl {
        if(root.children.length<=0) throw new Error("SRL.parseXml: xml is empty");

        const frameNodes = Array.from(root.querySelectorAll("frame"));
        // handle the case when there are no frames at all
        if(frameNodes.length == 0) {
            const tokens = root.textContent!.trim().split(/\s+/);
            return new Srl(tokens, []);
        }

        let sentence:string[] = [];

        const frames:Frame[] = [];
        for(const frameNode of frameNodes) {
            let tokenIndex = 0;
            const roles:Role[] = [];

            let iSentence:string[] = [];
            for(const childNode of Array.from(frameNode.childNodes)) {
                const text = childNode.textContent!.trim();
                const tokens = text.length>0?text.split(/\s+/):[];

                iSentence = [...iSentence, ...tokens];

                if(childNode.nodeType === Node.ELEMENT_NODE) {
                    const child = childNode as Element;
                    const roleLength = tokens.length;
                    const roleName = child.getAttribute('type') || "";
                    const isPredicate = child.tagName === 'pred';
                    roles.push(new Role(tokenIndex, roleLength, roleName, isPredicate));
                    tokenIndex += tokens.length;
                }
                else {
                    tokenIndex += tokens.length;
                }
            }
            frames.push(new Frame(roles));

            if(sentence.length == 0) {
                sentence = iSentence;
            }
            else if(sentence.length != iSentence.length) {
                throw new Error("Sentences in SRL frames are different.");
            }
        }

        return new Srl(sentence, frames, "");
    }

    toXml(document:Document):Element {
        const ret = document.createElement('sent');

        function appendToken(parent:Element, token:string) {
            if(parent.childNodes.length == 0) {
                parent.appendChild(document.createTextNode(token));
            }
            else {
                parent.appendChild(document.createTextNode(" " + token));
            }
        }

        function appendNode(parent:Element, node:Element) {
            if(parent.childNodes.length == 0) {
                parent.appendChild(node);
            }
            else {
                parent.appendChild(document.createTextNode(" "));
                parent.appendChild(node);
            }
        }

        if(this.frames.length === 0) {
            ret.textContent = this.tokens.join(" ");
            return ret;
        }

        ret.append(...this.frames.map(frame=>{
            const frameNode = document.createElement('frame');
            for(let i=0; i<this.tokens.length; i++) {
                const foundRole = frame.roles.find(role=>role.beginTokenIndex === i);
                if(foundRole == null) {
                    appendToken(frameNode, this.tokens[i]);
                    continue;
                }

                const roleNode = document.createElement(foundRole.isPredicate?"pred":"arg");
                roleNode.setAttribute("type", foundRole.name);

                for(let j=0; j<foundRole.length; j++) {
                    appendToken(roleNode, this.tokens[i + j]);
                }
                appendNode(frameNode, roleNode);
                i += foundRole.length - 1;
            }
            return frameNode;
        }));
        return ret;
    }

    static parseJson(hyp: SrlJsonFormat):Srl {
        const tokens = hyp.tokens;
        const frames = hyp.edges
            .map(edge=>{
                const tails = edge[0];
                const head = edge[1][0];
                const pred = Role.parseHypergraph(hyp.nodes[head], true);

                // check if the predicate is the &GOAL; node
                if(pred.length >= tokens.length) {
                    return null;
                }

                const roles = tails.map(index=>Role.parseHypergraph(hyp.nodes[index], false));
                return new Frame([pred, ...roles]);
            }).filter(x=>!!x) as Frame[];
        return new Srl(tokens, frames, hyp.meta?hyp.meta.frame_label_set_name:undefined);
    }

    static parseJsonhWithMap(hyp: SrlJsonFormat):{srl:Srl,hypNodeIndex_frameRoleIndex:Map<number,[number, number][]>} {
        const tokens = hyp.tokens;
        const hypNodeIndex_frameRoleIndex:Map<number,[number, number][]> = new Map();
        const frames = hyp.edges
            .filter((e,i)=>i>0) // strip away the first edge since it represents the list of predicates which is not used
            .map((edge, frameIndex)=>{
                const tails = edge[0];
                const head = edge[1][0];

                // store which hypergraph node corresponds to which constructed role
                const hypNodeIndex_role:Map<number, Role> = new Map();

                const pred = Role.parseHypergraph(hyp.nodes[head], true);
                hypNodeIndex_role.set(head, pred);

                const roles = tails.map(index=>{
                    const role = Role.parseHypergraph(hyp.nodes[index], false);
                    hypNodeIndex_role.set(index, role);
                    return role;
                });

                const ret = new Frame([pred, ...roles]);

                for(const [hypNodeIndex, role] of Array.from(hypNodeIndex_role.entries())) {
                    const roleIndex = ret.roles.indexOf(role);
                    const existingRole = hypNodeIndex_frameRoleIndex.get(hypNodeIndex);
                    if(existingRole == null) {
                        hypNodeIndex_frameRoleIndex.set(hypNodeIndex, [[frameIndex, roleIndex]]);
                    }
                    else {
                        existingRole.push([frameIndex, roleIndex])
                    }
                }

                return ret;
            });
        return {srl:new Srl(tokens, frames, hyp.meta?hyp.meta.frame_label_set_name:undefined), hypNodeIndex_frameRoleIndex};
    }


    static toAssert(srl:Srl):string {
        const lines = [];
        for(const frame of srl.frames) {
            let lineTokens = ['0:'];
            let pivot = 0;
            for(const role of frame.roles) {
                lineTokens = [...lineTokens, ...srl.tokens.slice(pivot, role.beginTokenIndex)];
                lineTokens = [...lineTokens, '['+role.name];
                lineTokens = [...lineTokens, ...srl.tokens.slice(role.beginTokenIndex, role.beginTokenIndex+role.length), ']'];
                pivot = role.beginTokenIndex+role.length;
            }
            lines.push(lineTokens.join(' '));
        }
        return lines.join('\n');
    }

    static createEmptySrlFromTokens(tokens:string[], frameLabelSetName:string) {
        return new Srl(tokens, [], frameLabelSetName);
    }

    static getSentence(srl:null|Srl) {
        if(srl == null) return "";
        return srl.tokens.join(' ');
    }

    static addFrame(srl:Srl, beginTokenIndex:number, endTokenIndex: number, name: string):Srl {
        const role = new Role(beginTokenIndex, endTokenIndex - beginTokenIndex, name, true);
        const frame = new Frame([role]);
        return new Srl(srl.tokens, [...srl.frames, frame])
    }

    static addRole(srl:Srl, frameIndex:number, beginTokenIndex:number, endTokenIndex:number, name:string):Srl {
        const len = endTokenIndex - beginTokenIndex;
        if(frameIndex >= srl.frames.length) throw new Error("Srl.addRole: frameIndex out of boundary");
        const frames = [...srl.frames];
        const roles = [...frames[frameIndex].roles];

        for(let i=roles.length-1; i>=0; i--) {
            const role = roles[i];
            if(!Role.isDisjoint(role.beginTokenIndex, role.length, beginTokenIndex, len)) {
                roles.splice(i, 1);
            }
        }
        const role = new Role(beginTokenIndex, endTokenIndex - beginTokenIndex, name, false);
        roles.push(role);
        frames[frameIndex] = new Frame(roles);
        return new Srl(srl.tokens, frames);
    }

    static removeFrame(srl:Srl, frameIndex:number) {
        if(frameIndex >= srl.frames.length) throw new Error("Srl.removeFrame: frameIndex out of boundary");
        const frames = [...srl.frames];
        frames.splice(frameIndex, 1);
        return new Srl(srl.tokens, frames);
    }

    static removeRole(srl:Srl, frameIndex:number, roleIndex:number):Srl {
        if(frameIndex >= srl.frames.length) throw new Error("Srl.removeRole: frameIndex out of boundary");
        const roles = [...srl.frames[frameIndex].roles];
        if(roleIndex >= roles.length)  throw new Error("Srl.removeRole: roleIndex out of boundary");
        if(roles[roleIndex].isPredicate) {
            return Srl.removeFrame(srl, frameIndex);
        }
        roles.splice(roleIndex,1);
        const frame = new Frame(roles);
        const frames = [...srl.frames];
        frames[frameIndex] = frame;
        return new Srl(srl.tokens, frames);
    }

    static changeRole(srl:Srl, frameIndex:number, roleIndex:number, name:string):Srl {
        if(frameIndex >= srl.frames.length) throw new Error("Srl.changeRole: frameIndex out of boundary");
        const roles = [...srl.frames[frameIndex].roles];
        if(roleIndex >= roles.length)  throw new Error("Srl.changeRole: roleIndex out of boundary");
        const old:Role = roles[roleIndex];
        roles[roleIndex] = new Role(old.beginTokenIndex, old.length, name, old.isPredicate);
        const frame = new Frame(roles);
        const frames = [...srl.frames];
        frames[frameIndex] = frame;
        return new Srl(srl.tokens, frames);
    }
}

export default Srl;