import org.jgrapht.graph.DefaultWeightedEdge;
import org.jgrapht.graph.SimpleWeightedGraph;
import org.jgrapht.nio.Attribute;
import org.jgrapht.nio.AttributeType;
import org.jgrapht.nio.DefaultAttribute;
import org.jgrapht.nio.graphml.GraphMLExporter;
import org.jgrapht.nio.graphml.GraphMLExporter.AttributeCategory;

import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class GraphMLFileWriter {

    String filename;
    SimpleWeightedGraph graph;
    GraphWithMapData data;
    Map<String, Attribute> graphAttributes;

    public GraphMLFileWriter(String filename, GraphWithMapData data) {
        if(!filename.matches(".*\\.graphml")){
            filename = filename + ".graphml";
        }
        this.filename = filename;
        this.data = data;
        this.graph = data.getGraph();
        graphAttributes = createGraphAttributes();
    }

    public GraphMLFileWriter(String filename, SimpleWeightedGraph<Vertex, DefaultWeightedEdge> graph) {
        if(!filename.matches(".*\\.graphml")){
            filename = filename + ".graphml";
        }
        this.filename = filename;
        this.graph = graph;
    }

    private Map<String, Attribute> createGraphAttributes(){
        Map<String, Attribute> attributes = new HashMap<>();
        //Sample plate filename
        attributes.put("sample plate filename", DefaultAttribute.createAttribute(data.getPlateFilename()));
        // Number of wells
        attributes.put("well count", DefaultAttribute.createAttribute(data.getNumWells().toString()));
        //Well populations
        Integer[] wellPopulations = data.getWellPopulations();
        StringBuilder populationsStringBuilder = new StringBuilder();
        populationsStringBuilder.append(wellPopulations[0].toString());
        for(int i = 1; i < wellPopulations.length; i++){
            populationsStringBuilder.append(", ");
            populationsStringBuilder.append(wellPopulations[i].toString());
        }
        String wellPopulationsString = populationsStringBuilder.toString();
        attributes.put("well populations", DefaultAttribute.createAttribute(wellPopulationsString));
        attributes.put("read depth", DefaultAttribute.createAttribute(data.getReadDepth().toString()));
        attributes.put("read error rate", DefaultAttribute.createAttribute(data.getReadErrorRate().toString()));
        attributes.put("error collision rate", DefaultAttribute.createAttribute(data.getErrorCollisionRate().toString()));
        attributes.put("real sequence collision rate", DefaultAttribute.createAttribute(data.getRealSequenceCollisionRate()));
        return attributes;
    }

    private Map<String, Attribute> createVertexAttributes(Vertex v){
        Map<String, Attribute> attributes = new HashMap<>();
        //sequence type
        attributes.put("type", DefaultAttribute.createAttribute(v.getType().name()));
        //sequence
        attributes.put("sequence", DefaultAttribute.createAttribute(v.getSequence()));
        //number of wells the sequence appears in
        attributes.put("occupancy", DefaultAttribute.createAttribute(v.getOccupancy()));
        //total number of times the sequence was read
        attributes.put("total read count", DefaultAttribute.createAttribute(v.getReadCount()));
        StringBuilder wellsAndReadCountsBuilder = new StringBuilder();
        Iterator<Map.Entry<Integer, Integer>> wellOccupancies = v.getWellOccupancies().entrySet().iterator();
        while (wellOccupancies.hasNext()) {
            Map.Entry<Integer, Integer> entry = wellOccupancies.next();
            wellsAndReadCountsBuilder.append(entry.getKey() + ":" + entry.getValue());
            if (wellOccupancies.hasNext()) {
                wellsAndReadCountsBuilder.append(", ");
            }
        }
        String wellsAndReadCounts = wellsAndReadCountsBuilder.toString();
        //the wells the sequence appears in and the read counts in those wells
        attributes.put("wells:read counts", DefaultAttribute.createAttribute(wellsAndReadCounts));
        return attributes;
    }

    public void writeGraphToFile() {
        try(BufferedWriter writer = Files.newBufferedWriter(Path.of(filename), StandardOpenOption.CREATE_NEW);
        ){
            //create exporter. Let the vertex labels be the unique ids for the vertices
            GraphMLExporter<Vertex, SimpleWeightedGraph<Vertex, DefaultWeightedEdge>> exporter = new GraphMLExporter<>(v -> v.getVertexLabel().toString());
            //set to export weights
            exporter.setExportEdgeWeights(true);
            //Set graph attributes
            exporter.setGraphAttributeProvider( () -> graphAttributes);
            //set type, sequence, and occupancy attributes for each vertex
            exporter.setVertexAttributeProvider(this::createVertexAttributes);
            //register the attributes
            for(String s : graphAttributes.keySet()) {
                exporter.registerAttribute(s, AttributeCategory.GRAPH, AttributeType.STRING);
            }
            exporter.registerAttribute("type", AttributeCategory.NODE, AttributeType.STRING);
            exporter.registerAttribute("sequence", AttributeCategory.NODE, AttributeType.STRING);
            exporter.registerAttribute("occupancy", AttributeCategory.NODE, AttributeType.STRING);
            exporter.registerAttribute("total read count", AttributeCategory.NODE, AttributeType.STRING);
            exporter.registerAttribute("wells:read counts", AttributeCategory.NODE, AttributeType.STRING);
            //export the graph
            exporter.exportGraph(graph, writer);
        } catch(IOException ex){
            System.out.println("Could not make new file named "+filename);
            System.err.println(ex);
        }
    }
}