package edu.berkeley.nlp.starcraft.construction;

import java.util.HashSet;
import java.util.Set;

import org.bwapi.proxy.model.Race;
import org.bwapi.proxy.model.TechType;
import org.bwapi.proxy.model.UnitType;
import org.bwapi.proxy.model.UpgradeType;

import edu.berkeley.nlp.starcraft.util.Digraph;
import edu.berkeley.nlp.starcraft.util.Log;

public class TechTree {

  public static TechTree TERRAN_TECH_TREE = new TechTree();
  public static TechTree PROTOSS_TECH_TREE = new TechTree();
  public static TechTree ZERG_TECH_TREE = new TechTree();

  public static TechTree getTechTreeFromRace(Race race) {
    TechTree tree = null;
    if (race.equals(Race.TERRAN)) {
      tree = TERRAN_TECH_TREE;
    } else if (race.equals(Race.ZERG)) {
      tree = ZERG_TECH_TREE;
    } else if (race.equals(Race.PROTOSS)) {
      tree = PROTOSS_TECH_TREE;
    } else {
      tree = new TechTree(); // Garbage tree, whatever.
    }
    return tree;
  }

  public Set<Node> requiredNodesFor(UnitType u) {
    return requiredNodesFor(new Unit(u));
  }

  public Set<Node> requiredNodesFor(TechType u) {
    return requiredNodesFor(new Tech(u));
  }

  public Set<Node> requiredNodesFor(UpgradeType u) {
    return requiredNodesFor(new Upgrade(u));
  }

  public Set<Node> requiredNodesFor(Node n) {
    Set<Node> ret = new HashSet<Node>();
    for (Digraph<Node>.Node nn : graph.getNode(n).incomingNodes) {
      ret.add(nn.data);
    }
    return ret;
  }
  
  
  /**
   * Returns a digraph representing the unit-type dependencies up to and
   * including the desired unit types but excludes unit types that exist.
   */
  public  Digraph<TechTree.Node> getBuildPath(Set<TechTree.Node> desiredUnitTypes,
      Set<TechTree.Node> existingUnitTypes) {

    return getGraph().subgraphWithSinksAndBlocks(desiredUnitTypes, existingUnitTypes);

  }

  public Set<Node> childNodesOf(Node n) {
    Set<Node> ret = new HashSet<Node>();
    for (Digraph<Node>.Node nn : graph.getNode(n).outgoingNodes) {
      ret.add(nn.data);
    }
    return ret;
  }

  // Only 3 Techtrees.
  private TechTree() {

  }

  private Digraph<Node> graph = new Digraph<Node>();

  public Digraph<Node> getGraph() {
    return graph;
  }

  public interface Node {
  }

  public static class Unit implements Node {
    @Override
    public String toString() {
      return "\"Unit (tpe=" + tpe.getName() + ")\"";
    }

    UnitType tpe;

    public UnitType getUnitType() {
      return tpe;
    }

    public Unit(UnitType tpe) {
      this.tpe = tpe;
    }

    @Override
    public boolean equals(Object obj) {
      if (obj == null) {
        return false;
      }
      if (getClass() != obj.getClass()) {
        return false;
      }
      final Unit other = (Unit) obj;
      if (this.tpe != other.tpe
          && (this.tpe == null || !this.tpe.equals(other.tpe))) {
        return false;
      }
      return true;
    }

    @Override
    public int hashCode() {
      int hash = 5;
      hash = 29 * hash + (this.tpe != null ? this.tpe.hashCode() : 0);
      return hash;
    }

  }

  public static class Tech implements Node {
    @Override
    public String toString() {
      return "\"Tech (tpe=" + tpe.getName() + ")\"";
    }

    TechType tpe;

    public TechType getTechType() {
      return tpe;
    }

    public Tech(TechType tpe) {
      this.tpe = tpe;
    }

    @Override
    public boolean equals(Object obj) {
      if (obj == null) {
        return false;
      }
      if (getClass() != obj.getClass()) {
        return false;
      }
      final Tech other = (Tech) obj;
      if (this.tpe != other.tpe
          && (this.tpe == null || !this.tpe.equals(other.tpe))) {
        return false;
      }
      return true;
    }

    @Override
    public int hashCode() {
      int hash = 7;
      hash = 37 * hash + (this.tpe != null ? this.tpe.hashCode() : 0);
      return hash;
    }

  }

  public static class Upgrade implements Node {
    @Override
    public String toString() {
      return "\"Upgrade (tpe=" + tpe.getName() + ")\"";
    }

    UpgradeType tpe;

    public UpgradeType getUpgradeType() {
      return tpe;
    }

    @Override
    public boolean equals(Object obj) {
      if (obj == null) {
        return false;
      }
      if (getClass() != obj.getClass()) {
        return false;
      }
      final Upgrade other = (Upgrade) obj;
      if (this.tpe != other.tpe
          && (this.tpe == null || !this.tpe.equals(other.tpe))) {
        return false;
      }
      return true;
    }

    @Override
    public int hashCode() {
      int hash = 7;
      hash = 67 * hash + (this.tpe != null ? this.tpe.hashCode() : 0);
      return hash;
    }

    public Upgrade(UpgradeType tpe) {
      this.tpe = tpe;
    }

  }

  private Node registerUnit(UnitType u) {
    Node n = new Unit(u);
    if (graph.containsNodeFor(n))
      return n;
    graph.addNode(n);
    for (UnitType r : u.requiredUnits().keySet()) {
      Node rn = registerUnit(r);
      graph.addEdge(rn, n);
    }
    TechType tech = u.requiredTech();
    if(!tech.equals(TechType.NONE)) {
      Node rn = registerTech(tech);
      graph.addEdge(rn, n);
    }
    return n;
  }

  private Node registerTech(TechType tech) {
    Node n = new Tech(tech);
    if (graph.containsNodeFor(n))
      return n;
    graph.addNode(n);

    UnitType u = tech.whatResearches();
    Node rn = registerUnit(u);
    graph.addEdge(rn, n);

    return n;
  }

  private Node registerUpgrade(UpgradeType ug) {
    Node n = new Upgrade(ug);
    if (graph.containsNodeFor(n))
      return n;
    graph.addNode(n);
   
    
    UnitType u = ug.whatUpgrades();
    Node rn = registerUnit(u);
    graph.addEdge(rn, n);
    return n;
  }
  
  @Override
  public String toString() {
    return "TechTree [graph=" + graph + "]";
  }


  static {
    TERRAN_TECH_TREE = new TechTree();
    PROTOSS_TECH_TREE = new TechTree();
    ZERG_TECH_TREE = new TechTree();
    for (UnitType u : UnitType.allUnitTypes()) {
      getTechTreeFromRace(u.getRace()).registerUnit(u);
    }

    for (TechType t : TechType.allTechTypes()) {
      Race r = t.whatResearches().getRace();
      getTechTreeFromRace(r).registerTech(t);
    }

    for (UpgradeType ug : UpgradeType.allUpgradeTypes()) {
      Race r = ug.whatUpgrades().getRace();
      getTechTreeFromRace(r).registerUpgrade(ug);
    }

  }
}
