package org.sparx.sequoia;

import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntList;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;

/**
 * Created by timbrooks on 6/16/15.
 */
public class Parser<F, C> {

    private final FeatureParser<F> featureParser;
    private List<IntList> categoricalSplits = new ArrayList<>();
    private Int2ObjectMap<F> featureIndexMap;

    public Parser(FeatureParser<F> featureParser) {
        this.featureParser = featureParser;
    }

    public Forest<Integer, double[]> parseForestFromStream(InputStream inputStream) throws IOException {
        InputStreamReader intputStream = new InputStreamReader(inputStream);
        BufferedReader reader = new BufferedReader(intputStream);

        ignoreHeader(reader);
        F[] features = featureParser.parseFeatures(reader);
        featureIndexMap = buildFeatureIndexMap(features);
        categoricalSplits = parseCategoricalLists(reader);
        List<Tree<Integer, double[]>> trees = parseTrees(reader);


        return Planter.createForestFromTrees(trees);
    }

    private List<Tree<Integer, double[]>> parseTrees(BufferedReader reader) throws IOException {
        int numberOfTrees = readNumberOfElements(reader.readLine());
        if (numberOfTrees <= 0) {
            throw new RuntimeException("numTrees has to be > 0");
        }
        List<Tree<Integer, double[]>> trees = new ArrayList<>();
        for (int i = 0; i < numberOfTrees; i++) {
            int numNodes = readNumberOfNodes(reader.readLine());

            List<Node<Integer, double[]>> nodes = new ArrayList<>();
            for (int j = 0; j < numNodes; j++) {
                nodes.add(parseNode(reader));
            }

            if (numNodes != nodes.size()) {
                String message = String.format("Number of nodes does not match expected: [expected: %s, actual: %s].",
                        numNodes, nodes.size());
                throw new RuntimeException(message);
            }
            trees.add(Planter.createTreeFromNodes(nodes));
        }
        return trees;
    }

    private Int2ObjectMap<F> buildFeatureIndexMap(F[] featuresFromModelFile) {
        Int2ObjectMap<F> result = new Int2ObjectOpenHashMap<>(featuresFromModelFile.length);
        for (int i = 0; i < featuresFromModelFile.length; i++) {
            result.put(i, featuresFromModelFile[i]);
        }
        return result;
    }


    private Node<Integer, double[]> parseNode(final BufferedReader reader) throws IOException {
        String[] nodeDataTokens = parseNodeData(reader.readLine());
        int featureIndex = Integer.parseInt(nodeDataTokens[0]);
        double cutPoint = Double.parseDouble(nodeDataTokens[1]);
        int[] childNodes = {Integer.parseInt(nodeDataTokens[2]),
                Integer.parseInt(nodeDataTokens[3])};
        double value = Double.parseDouble(nodeDataTokens[5]);

        boolean isLeaf = featureIndex == -1;
        if (isLeaf) {
            return new Node<>(featureIndex, value, isLeaf, childNodes, new LeafCondition());
        } else {
            return new Node<>(featureIndex, cutPoint, isLeaf, childNodes, new NumericCondition(cutPoint));
        }
    }

    private String[] parseNodeData(String line) {
        String[] nodeData = line.split(",");

        if (nodeData.length != 6) {
            throw new RuntimeException("expected 6 fields to specify a node");
        }
        for (int i = 0; i < nodeData.length; i++) {
            nodeData[i] = nodeData[i].trim();
        }
        return nodeData;
    }

    private static int readNumberOfNodes(String line) {
        if (!"tree [0-9]+".matches(line)) {
            throw new RuntimeException("Tree does not match pattern");
        }
        String[] tokens = line.split(" ");
        return Integer.parseInt(tokens[1].trim());
    }

    private List<IntList> parseCategoricalLists(BufferedReader reader) throws IOException {
        int numCategoricalLists = readNumberOfElements(reader.readLine());
        List<IntList> categoricalSplits = new ArrayList<>(numCategoricalLists);
        for (int i = 0; i < numCategoricalLists; i++) {
            categoricalSplits.add(categoricalSplitListFromLine(reader.readLine()));
        }

        return categoricalSplits;
    }

    private int readNumberOfElements(String line) {
        return Integer.parseInt(readKeyVal(line));
    }

    private void ignoreHeader(BufferedReader reader) throws IOException {
        reader.readLine();
    }

    private IntList categoricalSplitListFromLine(String line) {
        String[] categoricalIndexTokens = line.split(",");
        if (categoricalIndexTokens.length <= 0) {
            throw new RuntimeException("expected non-empty cat split indexes");
        }
        IntList result = new IntArrayList(categoricalIndexTokens.length);
        for (String categoricalIndexToken : categoricalIndexTokens) {
            result.add(Integer.parseInt(categoricalIndexToken.trim()));
        }
        return result;
    }

    private String readKeyVal(String line) {
        String[] tokens = line.split("=", 2);
        if (tokens.length != 2) {
            throw new RuntimeException(String.format(
                    "expected line in \"key = val\", got %s", line));
        }
        return tokens[1].trim();
    }
}

@@ -130,99 +130,8 @@
 Map;
-%0Aimport it.unimi.dsi.fastutil.ints.IntArrayList;%0Aimport it.unimi.dsi.fastutil.ints.IntList;
 %0A%0Aim
@@ -428,73 +428,8 @@
 er;%0A
-    private List%3CIntList%3E categoricalSplits = new ArrayList%3C%3E();%0A
     
@@ -691,85 +691,8 @@
     
-InputStreamReader intputStream = new InputStreamReader(inputStream);%0A        
 Buff
@@ -719,24 +719,46 @@
 new Buffered
+Reader(new InputStream
 Reader(intpu
@@ -754,17 +754,16 @@
 eader(in
-t
 putStrea
@@ -764,16 +764,17 @@
 tStream)
+)
 ;%0A%0A     
@@ -920,67 +920,8 @@
 s);%0A
-        categoricalSplits = parseCategoricalLists(reader);%0A
     
@@ -3823,1163 +3823,205 @@
 ate 
-List%3CIntList%3E parseCategoricalLists(BufferedReader reader) throws IOException %7B%0A        int numCategoricalLists = readNumberOfElements(reader.readLine());%0A        List%3CIntList%3E categoricalSplits = new ArrayList%3C%3E(numCategoricalLists);%0A        for (int i = 0; i %3C numCategoricalLists; i++) %7B%0A            categoricalSplits.add(categoricalSplitListFromLine(reader.readLine()));%0A        %7D%0A%0A        return categoricalSplits;%0A    %7D%0A%0A    private int readNumberOfElements(String line) %7B%0A        return Integer.parseInt(readKeyVal(line));%0A    %7D%0A%0A    private void ignoreHeader(BufferedReader reader) throws IOException %7B%0A        reader.readLine();%0A    %7D%0A%0A    private IntList categoricalSplitListFromLine(String line) %7B%0A        String%5B%5D categoricalIndexTokens = line.split(%22,%22);%0A        if (categoricalIndexTokens.length %3C= 0) %7B%0A            throw new RuntimeException(%22expected non-empty cat split indexes%22);%0A        %7D%0A        IntList result = new IntArrayList(categoricalIndexTokens.length);%0A        for (String categoricalIndexToken : categoricalIndexTokens) %7B%0A            result.add(Integer.parseInt(categoricalIndexToken.trim()));%0A        %7D%0A        return result
+int readNumberOfElements(String line) %7B%0A        return Integer.parseInt(readKeyVal(line));%0A    %7D%0A%0A    private void ignoreHeader(BufferedReader reader) throws IOException %7B%0A        reader.readLine()
 ;%0A  

Remove categorical code from parser