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.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("1EHZ");
            Structure right = reader.getStructureById("1EVV");

            List<Atom> listLeft = getPhosphorus(left.getChainByPDB("A"));
            List<Atom> listRight = getPhosphorus(right.getChainByPDB("A"));

            assert listLeft.size() == listRight.size();
            double rmsd = 0;
            for (int i = 0; i < listLeft.size(); i++) {
                Atom atomLeft = listLeft.get(i);
                Atom atomRight = listRight.get(i);

                double distance = calculateDistance(atomLeft, atomRight);
                rmsd += distance * distance;
            }
            rmsd /= listLeft.size();
            rmsd = Math.sqrt(rmsd);
            System.out.println(rmsd);
        } catch (IOException | StructureException e) {
            e.printStackTrace();
        }
    }

    private static List<Atom> getPhosphorus(Chain chain) {
        List<Atom> list = new ArrayList<>();

        for (Group group : chain.getAtomGroups()) {
            try {
                Atom atom = group.getAtom("P");
                list.add(atom);
            } catch (StructureException e) {
                System.err.println(e.getMessage());
            }
        }
        return list;
    }

    private static double calculateDistance(Atom left, Atom right) {
        double dx = left.getX() - right.getX();
        double dy = left.getY() - right.getY();
        double dz = left.getZ() - right.getZ();
        return Math.sqrt(dx * dx + dy * dy + dz * dz);
    }
}