package inport.ConversionUtils;

import java.io.FileWriter;
import java.io.IOException;
import java.util.*;
import java.util.function.Function;

public class MZnFormat {

    static public void write2DArrayOfInt(FileWriter writer,
                                          String name,
                                          ArrayList<ArrayList<Integer>> operations) throws IOException {
        locWrite2DArray(writer, name, operations, Object::toString, false);
    }

    static public void write2DArrayOfInt(FileWriter writer,
                                          String name,
                                          ArrayList<ArrayList<Integer>> operations,
                                          boolean isNumberedFromZero) throws IOException {
        locWrite2DArray(writer, name, operations, Object::toString, isNumberedFromZero);
    }

    static public void write2DArrayOfSet(FileWriter writer,
                                          String name,
                                          ArrayList<ArrayList<ArrayList<Integer>>> operations) throws IOException {
        locWrite2DArray(writer, name, operations,
                (ArrayList<Integer> a) -> {
                    StringBuilder s = new StringBuilder();
                    s.append('{');
                    boolean isFirst2 = true;
                    for (Integer val : a) {
                        if (isFirst2) {
                            isFirst2 = false;
                        } else {
                            s.append(", ");
                        }
                        s.append(val);
                    }
                    s.append('}');
                    return s.toString();
                }
        );
    }

    static public <T> void locWrite2DArray(FileWriter writer,
                                            String name,
                                            ArrayList<ArrayList<T>> operations,
                                            Function<T, String> toMZNFormat) throws IOException {
        locWrite2DArray(writer, name, operations, toMZNFormat, false);
    }

    static public int boolToInt(boolean b) {
        return b ? 1 : 0;
    }

    static public <T> void locWrite2DArray(FileWriter writer,
                                            String name,
                                            ArrayList<ArrayList<T>> operations,
                                            Function<T, String> toMZNFormat,
                                            boolean isNumberedFromZero) throws IOException {

        writer.write(name + " = array2d(" + (isNumberedFromZero ? 0 : 1) + ".." + (operations.size() - boolToInt(isNumberedFromZero))
                + ", " + (isNumberedFromZero ? 0 : 1) + ".." + (operations.get(0).size() - boolToInt(isNumberedFromZero))
                + ", [");

        boolean isFirst0 = true;

        for (ArrayList<T> a : operations) {
            for (T val : a) {
                if (isFirst0) {
                    isFirst0 = false;
                } else {
                    writer.write(", ");
                }
                writer.write(toMZNFormat.apply(val));
            }
        }
        writer.write("]);\n");
    }

    static public <T extends Collection<Integer>> void write2DArrayOfSetAs3DArray(FileWriter writer,
                                                                                   String name,
                                                                                   ArrayList<ArrayList<T>> operations) throws IOException {
        write2DArrayOfSetAs3DArray(writer, name, operations, false);
    }

    static public <T extends Collection<Integer>> void write2DArrayOfSetAs3DArray(FileWriter writer,
                                                                                   String name,
                                                                                   ArrayList<ArrayList<T>> operations,
                                                                                   boolean isNumberedFromZero) throws IOException {
        int maxSize = 0;
        ArrayList<ArrayList<Integer>> sizes = new ArrayList<>();

        int dim1 = operations.size();
        int dim2 = 0;

        for (ArrayList<T> a1 : operations) {
            dim2 = a1.size();
            ArrayList<Integer> locSizes = new ArrayList<>();
            for (T a2 : a1) {
                maxSize = Math.max(maxSize, a2.size());
                locSizes.add(a2.size());
            }
            sizes.add(locSizes);
        }
        writer.write(name + "_max_size = " + maxSize + ";\n");
        write2DArrayOfInt(writer, name + "_sizes", sizes, true);

        {
            String firstEl = Integer.toString(isNumberedFromZero ? 0 : 1);
            int modifier = isNumberedFromZero ? -1 : 0;
            writer.write(name + " = array3d(" + firstEl + ".." + (dim1 + modifier)
                    + ", " + firstEl + ".." + (dim2 + modifier)
                    + ", 1.." + name + "_max_size, [");
        }


        boolean isFirst = true;
        for (ArrayList<T> a1 : operations) {
            for (T a2 : a1) {
                for (Integer val : a2) {
                    if (isFirst) {
                        isFirst = false;
                    } else {
                        writer.write(", ");
                    }
                    writer.write(Integer.toString(val));
                }
                for (int i = a2.size(); i < maxSize; i++) {
                    if (isFirst) {
                        isFirst = false;
                    } else {
                        writer.write(", ");
                    }
                    writer.write(Integer.toString(1));
                }
            }
        }
        writer.write("]);\n");
    }

    static public <T> void write3dArray(FileWriter writer,
                                         String name,
                                         ArrayList<ArrayList<ArrayList<T>>> array,
                                         Function<T, String> toString) throws IOException {
        int dim1 = array.size();
        int dim2 = (dim1 > 0) ? array.get(0).size() : 0;
        int dim3 = (dim2 > 0) ? array.get(0).get(0).size() : 0;

        writer.write(name + " = array3d(1.." + dim1 + ", 1.." + dim2 + ", 1.." + dim3 + ", ");
        ArrayList<String> values = new ArrayList<>();
        for (ArrayList<ArrayList<T>> a2 : array) {
            for (ArrayList<T> a1 : a2) {
                for (T val : a1) {
                    values.add(toString.apply(val));
                }
            }
        }
        writer.write("[" + String.join(", ", values) + "]);\n");
    }

    static public void writeArrayOfSetAs2DArray(FileWriter writer,
                                                 String name,
                                                 ArrayList<ArrayList<Integer>> operations) throws IOException {
        writeArrayOfSetAs2DArray(writer, name, operations, true, false);
    }

    static public void writeArrayOfSetAs2DArray(FileWriter writer,
                                                 String name,
                                                 ArrayList<ArrayList<Integer>> operations,
                                                 boolean addAdditionalInfo,
                                                 boolean addDummyZeroElement) throws IOException {
        int maxSize = 0;
        ArrayList<Integer> sizes = new ArrayList<>();
        for (ArrayList<Integer> s : operations) {
            maxSize = Math.max(maxSize, s.size());
            sizes.add(s.size());
        }
        if (addAdditionalInfo) {
            writer.write(name + "_max_size = " + maxSize + ";\n");
            writeArray(writer, name + "_sizes", sizes);
        }
        writer.write(name + " = array2d(1.." + operations.size() + ", " + (addDummyZeroElement ? "0" : "1") + ".." + maxSize + ", [");

        boolean isFirst = true;
        for (ArrayList<Integer> s : operations) {
            if (addDummyZeroElement) {
                if (isFirst) {
                    isFirst = false;
                } else {
                    writer.write(", ");
                }
                writer.write("1");
            }
            for (Integer val : s) {
                if (isFirst) {
                    isFirst = false;
                } else {
                    writer.write(", ");
                }
                writer.write(Integer.toString(val));
            }
            for (int i = s.size(); i < maxSize; i++) {
                if (isFirst) {
                    isFirst = false;
                } else {
                    writer.write(", ");
                }
                writer.write("1");
            }
        }
        writer.write("]);\n");
    }

    static public <T, E> void writeArray(FileWriter writer,
                                          String name,
                                          ArrayList<T> data,
                                          Function<T, E> f,
                                          Optional<T> dummyZeroElement) throws IOException {
        writer.write(name + " = array1d(" + (dummyZeroElement.isPresent() ? 0 : 1) + ".." + data.size() + ", [");
        boolean isFirst = true;
        if (dummyZeroElement.isPresent()) {
            isFirst = false;
            writer.write(f.apply(dummyZeroElement.get()).toString());
        }
        for (int i = 0; i < data.size(); i++) {
            if (isFirst) {
                isFirst = false;
            } else {
                writer.write(", ");
            }
            writer.write(f.apply(data.get(i)).toString());
        }
        writer.write("]);\n");
    }

    static public <T> void writeArray(FileWriter writer, String name, ArrayList<T> data) throws IOException {
        writeArray(writer, name, data, (T val) -> val, Optional.empty());
    }

    static public <T, E> void writeArray(FileWriter writer,
                                          String name,
                                          ArrayList<T> data,
                                          Function<T, E> f) throws IOException {
        writeArray(writer, name, data, f, Optional.empty());
    }

    static public <T, E> void writeArray(FileWriter writer,
                                          String name,
                                          ArrayList<T> data,
                                          Optional<T> dummyZeroElement) throws IOException {
        writeArray(writer, name, data, (T val) -> val, dummyZeroElement);
    }

    static public <T> String setToString(Set<T> s) {
        StringBuilder res = new StringBuilder();
        res.append("{");
        boolean isFirst = true;
        for (T t : s) {
            if (isFirst) {
                isFirst = false;
            } else {
                res.append(", ");
            }
            res.append(t.toString());
        }
        res.append("}");
        return res.toString();
    }

    static public <T> String arrayToStringAsSet(ArrayList<T> a) {
        return setToString(new TreeSet<>(a));
    }
}