package etomica.modules.glass;

import etomica.atom.IAtom;
import etomica.box.Box;
import etomica.graphics.ColorScheme;
import etomica.space.Space;
import etomica.space.Vector;

import java.awt.*;

public class ColorSchemeDeviation extends ColorScheme {
    protected ConfigurationStorage configStorage;
    protected int configIndex;
    protected final Box box;
    protected final Vector dr;
    protected final Color[] colors;
    protected double fac;
    protected boolean isString;
    protected AtomTestDeviation atomTest;
    protected double dr2StringMax = 0.6*0.6;


    public ColorSchemeDeviation(Box box, ConfigurationStorage configStorage) {
        this.box = box;
        this.configStorage = configStorage;
        Space space = box.getSpace();
        dr = space.makeVector();
        colors = new Color[511];
        for (int i = 0; i < 256; i++) {
            colors[i] = new Color(0, i, 255 - i);
        }
        for (int i = 1; i < 256; i++) {
            colors[255 + i] = new Color(i, 255 - i, 0);
        }
        setLengthFactor(1);
        configIndex = 100;
    }

    public void setAtomTest(AtomTestDeviation atom) {atomTest = atom;}
    public void setIsString(boolean isString) {this.isString = isString;}
    public void setDrString(double drString) {this.dr2StringMax = drString*drString;}
    public double getDrString() {return Math.sqrt(this.dr2StringMax);}

    public boolean getIsString(){return isString;}

    public int getConfigIndex() {
        return configIndex;
    }

    public void setConfigIndex(int idx) {
        configIndex = idx;
    }

    public void setLengthFactor(double lenghtFactor) {
        this.fac = lenghtFactor;
    }

    public void setConfigStorage(ConfigurationStorage configStorage){
        this.configStorage = configStorage;
    }

    public ConfigurationStorage getConfigStorage(){return configStorage;}

    public double getDisplacementSq(IAtom a) {
        int idx = configIndex;
        int lastIndex = configStorage.getLastConfigIndex();
        if (idx > lastIndex) idx = lastIndex;
        if (idx == -1) return 0;
        Vector r = configStorage.getSavedConfig(0)[a.getLeafIndex()];
        Vector oldR = configStorage.getSavedConfig(idx)[a.getLeafIndex()];

        dr.Ev1Mv2(r, oldR);
        return dr.squared();
    }

    public double getRelativeDisplacement(IAtom a) {
        return Math.sqrt(getDisplacementSq(a)) / fac;
    }

    public boolean isAtomInString(IAtom atom){
        int idx = configIndex;
        int lastIndex = configStorage.getLastConfigIndex();
        if (idx > lastIndex) idx = lastIndex;
        if (idx == -1) return false;
        Vector r = configStorage.getSavedConfig(0)[atom.getLeafIndex()];
        for(int j=0;j<box.getLeafList().size();j++){
            Vector oldR = configStorage.getSavedConfig(idx)[j];
            dr.Ev1Mv2(r,oldR);
            box.getBoundary().nearestImage(dr);
            double dr2 = dr.squared();
            if(dr2 < dr2StringMax){
                return true;
            }
        }
        return false;
    }

    @Override
    public Color getAtomColor(IAtom a) {
        if(isString && isAtomInString(a)){
            return dr.getD()==3?Color.WHITE:Color.BLACK;
        }
        double s = getRelativeDisplacement(a);
        if (s >= 2) return colors[510];
        int i = (int) (510.9999 * s / 2);
        return colors[i];
    }
}

@@ -2731,19 +2731,25 @@
  for
+ 
 (int j
-=0;j%3C
+ = 0; j %3C 
 box.
@@ -2773,13 +2773,67 @@
 e();
+ 
 j++)
-%7B
+ %7B%0A            if (atom.getLeafIndex() == j) continue;
 %0A   

Not moving should not satisfy the string criteria