/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package etomica.modules.chainequilibrium;

import etomica.atom.AtomLeafAgentManager;
import etomica.atom.AtomType;
import etomica.atom.AtomTypeAgentManager;
import etomica.atom.IAtom;
import etomica.graphics.ColorScheme;
import etomica.simulation.Simulation;

import java.awt.*;
import java.util.Arrays;

/**
 * Color scheme for stepwise growth, based on the AtomType (alcohol vs. acid) and the number of bonds the atom can
 * form.
 *
 * @author Andrew Schultz
 */
public class ColorSchemeStepWise extends ColorScheme implements AtomTypeAgentManager.AgentSource {

    protected final AtomLeafAgentManager<IAtom[]> bondingAgentManager;
    protected final Simulation simulation;
    protected AtomTypeAgentManager[] colorMaps;

    public ColorSchemeStepWise(Simulation sim, AtomLeafAgentManager<IAtom[]> bondingAgentManager) {
        super();
        simulation = sim;
        colorMaps = new AtomTypeAgentManager[0];
        this.bondingAgentManager = bondingAgentManager;
    }

    public Color getAtomColor(IAtom atom) {
        IAtom[] nbrs = bondingAgentManager.getAgent(atom);
        if (nbrs != null && colorMaps.length > nbrs.length) {
            return (Color) colorMaps[nbrs.length].getAgent(atom.getType());
        }
        // we weren't told how to deal with any atom type with this many bonds.
        return ColorScheme.DEFAULT_ATOM_COLOR;
    }

    /**
     * Sets atoms of the given type and number of bonds to be the given color.
     */
    public void setColor(AtomType type, int nBonds, Color color) {
        if (nBonds >= colorMaps.length) {
            int oldLength = colorMaps.length;
            colorMaps = Arrays.copyOf(colorMaps, nBonds + 1);
            for (int i = oldLength; i < colorMaps.length; i++) {
                colorMaps[i] = new AtomTypeAgentManager(this, simulation);
            }
        }
        colorMaps[nBonds].setAgent(type, color);
    }

    public Object makeAgent(AtomType type) {
        return ColorScheme.DEFAULT_ATOM_COLOR;
    }

    public void releaseAgent(Object agent, AtomType type) {
    }
}

@@ -318,50 +318,8 @@
 pe;%0A
-import etomica.atom.AtomTypeAgentManager;%0A
 impo
@@ -420,27 +420,41 @@
 on;%0A
-%0A
 import 
-java.awt.*;
+etomica.species.ISpecies;%0A
 %0Aimp
@@ -466,19 +466,13 @@
 ava.
-util.Arrays
+awt.*
 ;%0A%0A/
@@ -688,52 +688,8 @@
 eme 
-implements AtomTypeAgentManager.AgentSource 
 %7B%0A%0A 
@@ -815,36 +815,23 @@
 otected 
-AtomTypeAgentManager
+Color%5B%5D
 %5B%5D color
@@ -992,46 +992,227 @@
     
-colorMaps = new AtomTypeAgentManager%5B0
+ISpecies lastSpecies = sim.getSpecies(sim.getSpeciesCount() - 1);%0A        AtomType lastAtomType = lastSpecies.getAtomType(lastSpecies.getAtomTypeCount() - 1);%0A        colorMaps = new Color%5BlastAtomType.getIndex() + 1%5D%5B4
 %5D;%0A 
@@ -1404,42 +1404,8 @@
 null
- && colorMaps.length %3E nbrs.length
 ) %7B%0A
@@ -1426,16 +1426,8 @@
 turn
- (Color)
  col
@@ -1437,45 +1437,47 @@
 aps%5B
-nbrs.length%5D.getAgent(atom.getType())
+atom.getType().getIndex()%5D%5Bnbrs.length%5D
 ;%0A  
@@ -1792,528 +1792,51 @@
     
-if (nBonds %3E= colorMaps.length) %7B%0A            int oldLength = colorMaps.length;%0A            colorMaps = Arrays.copyOf(colorMaps, nBonds + 1);%0A            for (int i = oldLength; i %3C colorMaps.length; i++) %7B%0A                colorMaps%5Bi%5D = new AtomTypeAgentManager(this, simulation);%0A            %7D%0A        %7D%0A        colorMaps%5BnBonds%5D.setAgent(type, color);%0A    %7D%0A%0A    public Object makeAgent(AtomType type) %7B%0A        return ColorScheme.DEFAULT_ATOM_COLOR;%0A    %7D%0A%0A    public void releaseAgent(Object agent, AtomType type) %7B
+colorMaps%5Btype.getIndex()%5D%5BnBonds%5D = color;
 %0A   

Simplify color scheme by dropping type agent manager (#106)