package pl.org.ptbi;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.biojava.bio.structure.Atom;
import org.biojava.bio.structure.Chain;
import org.biojava.bio.structure.Group;
import org.biojava.bio.structure.Structure;
import org.biojava.bio.structure.StructureException;
import org.biojava.bio.structure.align.BioJavaStructureAlignment;
import org.biojava.bio.structure.align.StructureAlignment;
import org.biojava.bio.structure.align.ce.CeMain;
import org.biojava.bio.structure.align.fatcat.FatCatFlexible;
import org.biojava.bio.structure.align.fatcat.FatCatRigid;
import org.biojava.bio.structure.align.model.AFPChain;
import org.biojava.bio.structure.align.seq.SmithWaterman3Daligner;
import org.biojava.bio.structure.io.PDBFileReader;

public class Main {
    public static void main(String[] args) {
        PDBFileReader reader = new PDBFileReader();
        reader.setAutoFetch(true);
        reader.setPath(System.getProperty("java.io.tmpdir"));

        try {
            Structure left = reader.getStructureById("3CFM");
            Structure right = reader.getStructureById("3UAF");

            Atom[] ca1 = getCarbonAlpha(left);
            Atom[] ca2 = getCarbonAlpha(right);

            StructureAlignment[] algorithms = new StructureAlignment[] { new CeMain(),
                    new SmithWaterman3Daligner(), new BioJavaStructureAlignment(),
                    new FatCatRigid(), new FatCatFlexible() };
            for (StructureAlignment algorithm : algorithms) {
                System.out.println(algorithm.getAlgorithmName());
                AFPChain alignment = algorithm.align(ca1, ca2);
                System.out.println(alignment.toCE(ca1, ca2));
            }
        } catch (IOException | StructureException e) {
            e.printStackTrace();
        }
    }

    private static Atom[] getCarbonAlpha(Structure structure) {
        List<Atom> list = new ArrayList<>();

        for (Chain chain : structure.getChains()) {
            for (Group group : chain.getAtomGroups()) {
                try {
                    Atom atom = group.getAtom("CA");
                    list.add(atom);
                } catch (StructureException e) {
                    System.err.println(e.getMessage());
                }
            }
        }
        return list.toArray(new Atom[list.size()]);
    }
}