package ca.bcgsc.abyssexplorer.graph;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.TreeMap;

import edu.uci.ics.jung.graph.DirectedSparseMultigraph;
import edu.uci.ics.jung.graph.util.EdgeType;

/**
 * Class to capture all the assembly information output by ABySS.
 * 
 * Built on the assumption that edge labels range from 0 -> n 
 * such that can use the label as an index
 * 
 * @author Cydney Nielsen
 * 
 */
public class AbyssGraph {
	
	protected String name;
	protected List<Vertex> adjVertices;
	protected Edge[] adjEdges;
	// map paired-end contig ids to member labels
	protected Map<Integer,ContigLabel[]> peMemberLabels;

	public AbyssGraph (int vCount, int eCount) {
		name = "";
		// only have an estimate of vertex count, so use and ArrayList
		Vertex[] v = new Vertex[vCount];
		adjVertices = new ArrayList<Vertex>(Arrays.asList(v));
		v = null; // dispose of initialization array
		// know exact edge count, so can use an Array
		adjEdges = new Edge[eCount];
	}
	
	public void setName(String name) {
		this.name = name;
	}
	
	public String getName() {
		return name;
	}
	
	// handle vertices
	
	public void setVertex(Vertex v) {
		int id = v.getId();
		if (id >= adjVertices.size()) {
			assert id == adjVertices.size(): "Error adding vertex";
			id = adjVertices.size();
			adjVertices.add(v);
		} else {
			adjVertices.set(id, v);
		}
	}
	
	public void setVertex(int index, Vertex v) throws IndexOutOfBoundsException {
		adjVertices.set(index, v);
	}
	
	public void clearVertex(Vertex v) {
		adjVertices.set(v.getId(), null);
	}
	
	public Vertex getVertex(int index) throws IndexOutOfBoundsException {
		return adjVertices.get(index);
	}
	
	public int getVertexCount() {
		return adjVertices.size();
	}
	
	// handle edges
	
	public void setEdge(Edge e) throws IndexOutOfBoundsException {
		setEdge(e.getId(), e);
	}
	
	public void setEdge(int index, Edge e) {
		adjEdges[index] = e;
	}
	
	/**
	 * Retrieves the edge with the corresponding label.
	 * 
	 * @param eLabel
	 * @return
	 * @throws IndexOutOfBoundsException
	 * @throws NumberFormatException
	 */
	public Edge getEdge(String eLabel) {
		
		int index = -1; int strand = 0;
		if (eLabel.endsWith("+")) {
			index = Integer.parseInt(eLabel.replace("+",""));
			strand = 0;
		} else if (eLabel.endsWith("-")) {
			index = Integer.parseInt(eLabel.replace("-",""));
			strand = 1;
		} else {
			throw new IllegalArgumentException("Invalid eLabel: '" + eLabel + "'");
		}
		Edge e = getEdge(index, strand);
		
		return e;
	}
	
	public Edge getEdge(ContigLabel l) {
		if (l == null) { return null; }
		return getEdge(l.getId(), l.getStrand());
	}
	
	/** 
	 * Retrieves the edge with the corresponding id and strand.
	 * More efficient than getEdge(String eLabel) if already have
	 * edge info in this form.
	 * 
	 * @param id
	 * @param strand
	 * @return
	 * @throws IndexOutOfBoundsException
	 */
	public Edge getEdge(int id, int strand) throws IndexOutOfBoundsException {
		Edge e = getEdge(id);
		// ensure the edge is oriented as requested
		if (e != null && e.getStrand() != strand) {
			e.reverseComplement();
		}
		return e;
	}
	
	/**
	 * Retrieves the edge with the corresponding id.
	 * 
	 * @param id
	 * @return
	 * @throws IndexOutOfBoundsException
	 */
	public Edge getEdge(int id) throws IndexOutOfBoundsException {
		return adjEdges[id];
	}
	
	public int getEdgeCount() {
		return adjEdges.length;
	}
	
	public Collection<Edge> getIncidentEdges(Vertex v) {
		Collection<Edge> incident = new HashSet<Edge>();
		List<Edge> in = v.getIncoming();
		if (in != null) { incident.addAll(in); }
		List<Edge> out = v.getOutgoing();
		if (out != null) { incident.addAll(out); }
        return incident;
	}
	
	public Collection<Edge> getIncidentEdges(Edge e) {
		Collection<Edge> incident = new HashSet<Edge>();
		for (Vertex v: e.getVertices()) {
			Collection<Edge> v_incident = getIncidentEdges(v);
			if (v_incident != null) { incident.addAll(v_incident); }
		}
		return incident;
	}

	public Collection<Vertex> getEndpoints(Edge e) {
		Collection<Vertex> endpoints = new ArrayList<Vertex>(2);
		endpoints.add(e.getSourceVertex());
		endpoints.add(e.getDestVertex());
		return endpoints;
	}
	
	public boolean hasEdge(String label) {
		if (getEdge(label) != null) {
			return true;
		} else {
			return false;
		}
	}
	
	public boolean hasPeContig(String peLabel) {
		ContigLabel l = new ContigLabel(peLabel);
		return peMemberLabels.containsKey(l.getId());
	}
	
	/**
	 * Returns an ordered list of member single-end contigs 
	 * for the requested paired-end contig id. Contig order 
	 * and strand are appropriate for the requested paired-end 
	 * contig strand.
	 * 
	 * @param peLabel
	 * @return
	 */
	public List<Edge> getPairedEndContigMembers(ContigLabel peLabel) {
		if (peMemberLabels == null) { return null; }
		ContigLabel[] mLabels;
		try {
			mLabels = peMemberLabels.get(peLabel.getId());
		} catch (NullPointerException e) {
			throw new IllegalArgumentException("No paired-end contig with id " + peLabel.getId());
		}
		List<Edge> members = new ArrayList<Edge>(mLabels.length);
		if (peLabel.getStrand() == 0) {
			// values as stored
			for (ContigLabel l: mLabels) {
				members.add(getEdge(l.getId(), l.getStrand()));
			}
		} else {
			// return reverse order and opposite strands
			for (int i=mLabels.length-1; i>-1; i--) {
				ContigLabel l = mLabels[i];
				if (l.getStrand() == 0) {
					members.add(getEdge(l.getId(), 1));
				} else {
					members.add(getEdge(l.getId(), 0));
				}
			}
		}
		return members;
	}
	
	public void setPairedEndContig(int peId, List<ContigLabel> members) throws IllegalArgumentException {
		// ensure no id clash with single-end contigs
		if (peId < adjEdges.length) {
			throw new IllegalArgumentException("Single-end/Paired-end contig id clash: " + peId);
		}
		if (peMemberLabels == null) {
			peMemberLabels = new HashMap<Integer,ContigLabel[]>();
		}
		// store this peId within each member edge and collect an array of members for storage
		ContigLabel m; Edge mEdge; 
		ContigLabel[] mLabels = new ContigLabel[members.size()]; // array form
		for (int i=0; i<members.size(); i++) {
			m = members.get(i);
			mEdge = getEdge(m.getId(), m.getStrand());
			if (mEdge == null) {
				throw new IllegalArgumentException("Unknown single-end contig (" + m + ")");
			}
			// store pe orientation relative to mEdge's positive strand
			if (mEdge.getStrand() == 0) {
				mEdge.addPairedEndId(new ContigLabel(peId,(byte) 0));
			} else {
				// store rc of pe
				mEdge.addPairedEndId(new ContigLabel(peId,(byte) 1));
			}
			mLabels[i] = m;
		}
		// override any existing values for this peId
		peMemberLabels.put(peId, mLabels);
	}
	
    public DirectedSparseMultigraph<Vertex,Edge> getQuerySubgraph(Edge e, int maxEdges) {
        
    	DirectedSparseMultigraph<Vertex,Edge> subGraph = new DirectedSparseMultigraph<Vertex,Edge>();           
    	Queue<Edge> queue = new LinkedList<Edge>();
    	Map<Integer,Integer> inQueue = new TreeMap<Integer,Integer>();
    	queue.offer(e);
    	inQueue.put(e.getId(), 1);

    	int totalEdges = 0;
    	while (queue.peek() != null) {
    		Edge cEdge = queue.poll();
    		subGraph.addEdge(cEdge, cEdge.getSourceVertex(), cEdge.getDestVertex(), EdgeType.DIRECTED);
    		totalEdges++;
    		if (totalEdges > maxEdges) { break; }

    		// add this neighbors to the queue
    		Collection<Edge> incident = getIncidentEdges(cEdge);
    		for (Edge iEdge: incident) {
    			if (iEdge.getId() == cEdge.getId()) { continue; } // skip self
    			if (inQueue.containsKey(iEdge.getId())) { continue; }
    			inQueue.put(iEdge.getId(), 1);
    			queue.offer(iEdge);
    		}
    	}
    	return subGraph;
    }

}