Gedcom2Dot.java

/*
 * Copyright 2011 Foundation for On-Line Genealogy, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.folg.gedcom.tools;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

import org.folg.gedcom.model.EventFact;
import org.folg.gedcom.model.Family;
import org.folg.gedcom.model.Gedcom;
import org.folg.gedcom.model.Name;
import org.folg.gedcom.model.Person;
import org.folg.gedcom.parser.ModelParser;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;
import org.kohsuke.args4j.Option;
import org.xml.sax.SAXParseException;
 

/**
 * Generates a GraphViz input file from a gedcom model starting with a root person id.
 * 
 * To run the tool, first run with the switch --list to find the root person id
 * 
 * The output file can be converted to a PDF, using commands like this:
 * dot  -Tpdf -o familytree.pdf generatedfile.dot
 * 
 * @author Dominik Seichter (2021)
 */
public class Gedcom2Dot {
	@Option(name = "-i", required = true, usage = "gedcom file in")
	private File gedcomIn;

	@Option(name = "-o", required = false, usage = "dot file out")
	private File dotOut;

	@Option(name = "-l", required = false, usage = "max level")
	private int maxLevel = -1;
	
	@Option(name = "--list", required = false, usage = "List all persons in the gedcom file")
	private boolean listMode = false;
	
	@Option(name = "-r", required = false, usage = "ID of root person")
	private String rootId;
	
	private int maxDepth = 0;
	
	private Set<String> visited = new HashSet<>();
	
	private Gedcom gedcom;

	private StringBuffer nodes = new StringBuffer();
	private StringBuffer edges = new StringBuffer();
	private StringBuffer subs = new StringBuffer();

	/**
	 * List all persons in the file to determine ids
	 * @throws IOException 
	 * @throws SAXParseException 
	 */
	public void list() throws SAXParseException, IOException {
		final ModelParser modelParser = new ModelParser();
		gedcom = modelParser.parseGedcom(gedcomIn);

		final List<Person> persons = gedcom.getPeople();
		for( final Person p : persons) {
			System.out.println(p.getId() + " " + getNameOfPerson(p));
		}
	}
	
	/** 
	 * Process a gedcom model to generate a dot file
	 * 
	 * @throws SAXParseException
	 * @throws IOException
	 */
	public void process() throws SAXParseException, IOException {
		final ModelParser modelParser = new ModelParser();
		gedcom = modelParser.parseGedcom(gedcomIn);
		gedcom.createIndexes();

		Person root = findPerson(rootId);
		if( root == null) {
			throw new IllegalArgumentException("Person with root id was not found: " + rootId);
		}
		walkTree(root, 0);

		if (dotOut != null) {
			PrintWriter writer = new PrintWriter(dotOut);
			writer.println("// Generated by Dominik Seichter.\n");
			writer.println("digraph");
			writer.println("{");
			// Direction bottom -> tob
			writer.println("rankdir = BT;");
			// Orthogonal lines instead of curved
			writer.println("splines = ortho;");
			writer.println("overlap = false;");
			
			writer.println("\n");
			writer.println(nodes.toString());
			writer.println(edges.toString());
			writer.println(subs.toString());
			writer.println("}");
			writer.close();
		}
		
		System.out.println("Maximum depth: " + maxDepth);
	}

	private void walkTree(Person person, int level) {
		if (person == null || 
				// Do not traverse further than max level
				level == maxLevel || 
				// Certain models contain same ID several times, so do not generate extra node
				visited.contains(person.getId())) {
			return;
		}

		maxDepth = Math.max(level, maxDepth);
		visited.add(person.getId());
		
		final String prefix = "-".repeat(level);
		final String name = getNameOfPerson(person);
		final String dates = getDatesOfPerson(person);
		final String label = name + "\\n" + dates;
			
		System.out.println(prefix + label);

		nodes.append(person.getId() + " [shape=box, label=\"" + label + "\", color=black ];\n");
		
		final List<Family> parentFamilies = person.getParentFamilies(gedcom);
		for (Family f : parentFamilies) {
			f.getHusbands(gedcom)
				.stream().forEach(h -> walkTree(h, level + 1));
			f.getWives(gedcom)
				.stream().forEach(w -> walkTree(w, level + 1));

			final List<String> husbandIds = f.getHusbands(gedcom)
					.stream().map(Person::getId).collect(Collectors.toList());
			final List<String> wiveIds = f.getWives(gedcom)
					.stream().map(Person::getId).collect(Collectors.toList());
			final List<String> allIds = new ArrayList<>();
			allIds.addAll(husbandIds);
			allIds.addAll(wiveIds);
			
			subs.append("subgraph fam_" + f.getId() + " {\n");
			subs.append("\t{rank=same " + String.join(" ", allIds) + "}\n");
			subs.append("}\n");
			
			f.getHusbands(gedcom)
				.stream().forEach(h -> edges.append(person.getId() + " -> " + h.getId() + " [dir=none];\n"));
			f.getWives(gedcom) 
				.stream().forEach(w -> edges.append(person.getId() + " -> " + w.getId() + " [dir=none];\n"));
		}
	}

	private String getDatesOfPerson(Person person) {
		final StringBuffer b = new StringBuffer();
		final Optional<EventFact> birthDate = person.getEventsFacts()
				.stream()
				.filter( ef -> "Birth".equals(ef.getDisplayType()) )
				.findFirst();
		
		final Optional<EventFact> deathDate = person.getEventsFacts()
				.stream()
				.filter( ef -> "Death".equals(ef.getDisplayType()) )
				.findFirst();
		
		if( birthDate.isPresent() && !isNull(birthDate.get().getDate()) ) {
			b.append(" *" + birthDate.get().getDate() );
			if( !isNull(birthDate.get().getPlace()) ) {
				b.append(" (" + birthDate.get().getPlace() + ")");
			} 
		}
		
		if( deathDate.isPresent() && !isNull(deathDate.get().getDate()) ) {
			if( b.length() > 0 ) {
				b.append("\\n");
			}
			
			b.append(" +" + deathDate.get().getDate() );
			
			if( !isNull(deathDate.get().getPlace()) ) {
				b.append(" (" + deathDate.get().getPlace() + ")");
			} 
		}
		
		
		return b.toString();
	}

	/**
	 * Check if str is null or equals "null"
	 * @param str
	 * @return
	 */
	private boolean isNull(String str) {
		return str == null || str.equals("null");
	}
	
	/**
	 * Get the main name of a person and escape it for GraphViz output
	 * @param person 
	 * @return display name of the person
	 */
	private String getNameOfPerson(Person person) {
		if (!person.getNames().isEmpty()) {
			Name n = person.getNames().get(0);
			return n.getDisplayValue().replace("\"", "\\\"");
		} else {
			return "<Unknown>";
		}
	}

	/**
	 * Retrieve person with id from the model
	 * @param id
	 * @return
	 */
	private Person findPerson(final String id) {
		return gedcom.getPerson(id);
	}
	
	/**
	 * Utiltiy method to fine person via name - unused - only for debugging
	 * @param given
	 * @param sur
	 * @return
	 */
	@SuppressWarnings("unused")
	private Person findPerson(final String given, final String sur) {
		for (Person p : gedcom.getPeople()) {
			Optional<Name> found = p.getNames().stream()
					.filter(n -> given.equals(n.getGiven()) && sur.equals(n.getSurname())).findFirst();
			if (found.isPresent()) {
				return p;
			}
		}

		return null;
	}

	public static void main(String[] args) throws SAXParseException, IOException {
		Gedcom2Dot self = new Gedcom2Dot();
		CmdLineParser parser = new CmdLineParser(self);
		try {
			parser.parseArgument(args);
			
			if( self.isListMode() ) {
				self.list();
			} else {
				self.process();
			}
			
		} catch (CmdLineException e) {
			System.err.println(e.getMessage());
			parser.printUsage(System.err);
		}
	}

	public File getGedcomIn() {
		return gedcomIn;
	}

	public void setGedcomIn(File gedcomIn) {
		this.gedcomIn = gedcomIn;
	}

	public File getDotOut() {
		return dotOut;
	}

	public void setDotOut(File dotOut) {
		this.dotOut = dotOut;
	}

	public int getMaxLevel() {
		return maxLevel;
	}

	public void setMaxLevel(int maxLevel) {
		this.maxLevel = maxLevel;
	}

	public String getRootId() {
		return rootId;
	}

	public void setRootId(String rootId) {
		this.rootId = rootId;
	}

	public boolean isListMode() {
		return listMode;
	}

	public void setListMode(boolean listMode) {
		this.listMode = listMode;
	}


}