package beast.evolution.tree;

import beast.core.BEASTObject;
import beast.core.Description;
import beast.core.Input;
import beast.core.Input.Validate;
import beast.core.Loggable;

import java.io.PrintStream;

@Description("Logs tree annotated with metadata and/or rates")
public class NodeStatesLogger extends BEASTObject implements Loggable {
    final public Input<NodeStatesArray> nodesStatesInput = new Input<>("nodesStates",
            "Internal node states to be logged", Validate.REQUIRED);
    final public Input<Tree> treeInput = new Input<>("tree",
            "tree to provide node ID while logging.", Validate.REQUIRED);

//    final public Input<Boolean> compressInput = new Input<>("compress",
//            "if require to compress log", false, Validate.OPTIONAL);

    private NodeStatesArray nsa;
    private int tipCount;
    private int nodeCount;

    private boolean isCompress;

    @Override
    public void initAndValidate() {
        nsa = nodesStatesInput.get();
        tipCount = nsa.getTipsCount();
        nodeCount = nsa.getNodeCount();

//        isCompress = compressInput.get();
    }

    @Override
    public void init(PrintStream out) {
        out.println("#Sequences\t" + nsa.getInternalNodeCount());
        out.println("#Sites\t" + nsa.getSiteCount());
        out.println("#States\t" + nsa.getStateCount());

        out.println("Sample\tNode\tStates");
        // state from 0 to 59/60

//        final int ncol = nsa.getSiteCount();
//        if (ncol == 1) {
//            out.print(getID() + "\t");
//        } else {
//            for (int i = 0; i < ncol; i++) {
//                out.print("s" + (i + 1) + "\t");
//            }
//        }
    }

    //TODO multithreading https://stackoverflow.com/questions/21632585/thread-safety-of-printstream-in-java
    @Override
    public void log(long sampleNr, PrintStream out) {
        NodeStatesArray nsaCurrent = (NodeStatesArray) nsa.getCurrent();
        Tree treeCurrent = (Tree) treeInput.get().getCurrent();
        int branch = treeCurrent.getInternalNodeCount() - 1;
        out.print(sampleNr + "\t" + branch + "\t");
        // log branches, parent node .. child, e.g. 3..1,5..3,3..2,5..4
        out.print(getNodesMap(treeCurrent));
        // only internal nodes
        for (int i=tipCount; i < nodeCount; i++) {
            out.print("\t");
            nsaCurrent.getNodeStates(i).log(out); // state from 0 to 59/60
            out.print("\n");
        }
    }

    // parent node .. child, e.g. 3..1,5..3,3..2,5..4
    private String getNodesMap(Tree tree) {
        StringBuilder nodesMap = new StringBuilder();
        for (int i = 0; i < tree.getNodeCount(); i++) {
            Node node = tree.getNode(i);
            if (!node.isRoot()) {
                Node parent = node.getParent();
                int pa = parent.getNr();
                // parent node .. child
                nodesMap.append(pa).append("..").append(i);

                if(i < (tree.getNodeCount() - 1))
                    nodesMap.append(",");
            }
        }
        return nodesMap.toString();
    }

    @Override
    public void close(PrintStream out) {
// nothing to do
    }

}


@@ -1190,16 +1190,30 @@
 intln(%22#
+Internal Node 
 Sequence
@@ -1240,32 +1240,32 @@
 alNodeCount());%0A
-
         out.prin
@@ -2022,69 +2022,8 @@
 ();%0A
-        int branch = treeCurrent.getInternalNodeCount() - 1;%0A
     
@@ -2054,22 +2054,17 @@
  %22%5Ct%22 + 
-branch
+0
  + %22%5Ct%22)
@@ -2178,24 +2178,49 @@
 eCurrent));%0A
+        out.print(%22%5Cn%22);%0A
         // o
@@ -2489,24 +2489,77 @@
 3,3..2,5..4%0A
+    // Note: BEAST tree log makes nodeNr+1 in Newick%0A
     private 
@@ -2938,16 +2938,18 @@
 ppend(pa
++1
 ).append
@@ -2963,16 +2963,18 @@
 append(i
++1
 );%0A%0A    

fix bug of new line