package pl.org.ptbi;

import java.io.IOException;

import org.biojava.bio.structure.Atom;
import org.biojava.bio.structure.Chain;
import org.biojava.bio.structure.Element;
import org.biojava.bio.structure.Group;
import org.biojava.bio.structure.Structure;
import org.biojava.bio.structure.io.PDBFileReader;

import com.google.common.collect.SortedMultiset;
import com.google.common.collect.TreeMultiset;

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

        Structure structure;
        try {
            structure = reader.getStructureById("1EHZ");
        } catch (IOException e) {
            e.printStackTrace();
            return;
        }

        SortedMultiset<Element> bag = TreeMultiset.create();
        double mass = 0;
        for (Chain chain : structure.getChains()) {
            for (Group group : chain.getAtomGroups()) {
                for (Atom atom : group.getAtoms()) {
                    Element element = atom.getElement();
                    bag.add(element);
                    mass += element.getAtomicMass();
                }
            }
        }

        System.out.println("Total mass: " + mass);
        System.out.println("Atom count: " + bag.size());
        System.out.println("Elements present:");
        for (Element element : bag.elementSet()) {
            System.out.println(element + " : " + bag.count(element));
        }
    }
}