package pdb_movie;

/**
 * <p>Title: structureNMRImpl</p>
 *
 * <p>Description: For handling Large PDB files</p>
 *
 * @author Horvath Tamas
 * @version 0.5
 */

import java.io.*;
import java.util.*;
import org.biojava.bio.structure.*;
import org.biojava.bio.structure.io.PDBFileParser;
import java.util.zip.Adler32;
import java.util.zip.GZIPOutputStream;
import java.util.zip.GZIPInputStream;
import java.util.zip.CheckedOutputStream;
import java.util.jar.*;

public class structureNMRImpl {
    private HashMap<Integer,Integer> fileAccessMap = new HashMap<Integer,Integer>();
    private Structure headerContainer = new StructureImpl();
    private Integer modelnr = 0;
    private LineNumberReader origPDB = null;
    private PDBFileParser pdbfileparser = new PDBFileParser();
    private Integer currentModelnr = 0;
    private HashMap<Integer,String> modifiedModelMap = new HashMap<Integer,String>();
    private Structure currentModel =  new StructureImpl();
    private String baseFileName = "";
    private String jarFileName = "";
    private String PDBFileName = "";


    public structureNMRImpl() {


    }

    public void setBaseFileName(String fn) {
        baseFileName = fn;
    }

    private void resetIS() {

        if (jarFileName != null && !jarFileName.matches("\\s*") &&
            PDBFileName != null && !PDBFileName.matches("\\s*")   )
              resetIS(jarFileName,PDBFileName,"Jar");
        //many more options can be added here (like Zip, GZIP,uncompressed)
        else System.err.println("Unable to reset file!");
    }

    private void resetIS(String jf,String je, String mode) {

        LineNumberReader buf = null;

        try {
            if (mode.equalsIgnoreCase("Jar")) {
                JarFile jarfile = new JarFile(jf);
                JarEntry jarentry = jarfile.getJarEntry(je);
                buf = new LineNumberReader(new InputStreamReader(jarfile.
                        getInputStream(jarentry)));
            }
            //again many more options can be put here according to resetIS()
            else System.err.println("Unrecognised Mode");

        } catch (Exception e) {
            e.printStackTrace();
            System.err.println("Error initializing LineNumberReader");
        }
        origPDB = buf;

    }

    public structureNMRImpl(String jf,String je,String mode) {
        if (mode.equalsIgnoreCase("Jar")) {
            jarFileName = jf;
            PDBFileName = je;
        }
        //more modes here
        else System.err.println("Unrecognised Mode");

        resetIS();
        readNMRStructure();
    }


    public void readNMRStructure() {
        System.out.println("Indexing.");
        indexPDB();
        System.out.println("Reading header.");
        readHeader();
        if (fileAccessMap != null) modelnr = fileAccessMap.size();

        currentModelnr = 0;
    }

    public Integer nrModels() {
        return modelnr;
    }
    //I use the following methods to get around the PDB, but maybe some sort of iterator could be created too
    public boolean hasNextModel() {
        if (modelnr != null && currentModelnr < modelnr - 1) return true;
        return false;
    }

    public Structure nextModel() {
        if(!hasNextModel()) return null;
        else return getModel(++currentModelnr);
    }

    //positioning methods
    private void skipLines(int pos) {
        try {
            for (int i = 0; i < pos; i++)
                origPDB.readLine();
        }
        catch (Exception e) {
            System.err.println("Error while seeking position");
            e.printStackTrace();
        }
    }

    private void seekModel(int pos) {

        int currentPosition = origPDB.getLineNumber();
        //System.out.println("Jumping from "+currentPosition+" to "+pos);
        if ( currentPosition > pos) {
            try {
                resetIS();
                skipLines(pos);
            }
            catch (Exception e) {
                System.err.println("Could not reset and skip to position");
                e.printStackTrace();
            }
        }
        else {
            try {
                skipLines(pos-currentPosition);
            }
            catch (Exception e) {
                System.err.println("Could not skip to position");
                e.printStackTrace();
            }

        }
    }


    public Structure getModel(int mnr) {
        ArrayList<String> model = new ArrayList<String>();
        saveModel();
        if (fileAccessMap.containsKey(mnr)){

            try {
                if (modifiedModelMap.get(currentModelnr) == null) {
                    seekModel(fileAccessMap.get(mnr));
                    while (origPDB.ready()) {
                        String line = origPDB.readLine();
                        if (line.startsWith("MODEL") || line.startsWith("ENDMDL"))
                            break;
                        model.add(line);
                    }
                    currentModel = pdbfileparser.parsePDBFile(model);
                } else {
                    //the modified models are saved as gzips (we want to spare space)
                    System.out.println("Reading modified model");
                    String fileName = modifiedModelMap.get(currentModelnr);
                    GZIPInputStream pdbgz = new GZIPInputStream(new FileInputStream(
                            fileName));
                    BufferedReader buf = new BufferedReader(new InputStreamReader(pdbgz));
                    currentModel = pdbfileparser.parsePDBFile(buf);
                }
                currentModelnr = mnr;
                return currentModel;
            } catch (Exception e) {
                System.err.println("Error reading Model nr.:" + mnr + "/" + modelnr);
                e.printStackTrace();
            }
        }
        return null;
    }

    private void readHeader() {
        ArrayList<String> headerLines =  new ArrayList<String>();
        try {
            resetIS();
            while(origPDB.ready()) {
                String line = origPDB.readLine();
                if (line.startsWith("MODEL")) break;
                headerLines.add(line);
            }

            headerContainer = pdbfileparser.parsePDBFile(headerLines);
        }
        catch (Exception e) {
            System.err.println("Error reading PDB Header");
            e.printStackTrace();
        }
    }

    /*
      Handler for
      MODEL Record Format

      COLUMNS       DATA TYPE      FIELD         DEFINITION
      ----------------------------------------------------------------------
      1 -  6       Record name    "MODEL "
      11 - 14       Integer        serial        Model serial number.
     */

    //this method indexes the PDB to enable easier access of the models
    private void indexPDB() {
        try {
            int fileIndex = 0;
            while (origPDB.ready()) {
                String line = origPDB.readLine();
                if (line.startsWith("MODEL")) {
                    String modelNrString = line.substring(10,14);
                    Integer modelnr = Integer.parseInt(modelNrString.trim());
                    fileIndex = origPDB.getLineNumber();
                    fileAccessMap.put(modelnr,fileIndex);
                    System.out.print("\rIndexing Model "+ modelnr +"@"+fileIndex);
                }
            }
        }
        catch (Exception e) {
            System.err.println("Error indexing PDB");
            e.printStackTrace();
        }
        System.out.println("\nDone");
    }

    private String createModelFileName() {
        String mnr = currentModelnr.toString();
        String modelnrstring = "";

        //can it be done more elegantly?
        for (int i = 0; i < 4-mnr.length(); i++) modelnrstring += "0";
        modelnrstring += mnr;
        return baseFileName+"_model"+modelnrstring;
    }

    //this method saves the current model to a temporary file in gzip format
    private void saveModel() {
        if (currentModel != null && currentModel.size() > 0) {
            String fileName = createModelFileName();
            try {
                FileOutputStream dest = new FileOutputStream(fileName);
                GZIPOutputStream out = new GZIPOutputStream(new PrintStream(
                        dest));

                PrintStream ps = new PrintStream(out);
                ps.print(currentModel.toPDB().trim());
                out.close();

            } catch (Exception e) {
                System.err.println("Error saving Model");
                e.printStackTrace();
            }

            modifiedModelMap.put(currentModelnr, fileName);
        }
    }

    public void addModel(Structure model, int modelnum) {
        currentModel = model;
        currentModelnr = modelnum;
        if (fileAccessMap.isEmpty()) { modelnum = 0; currentModelnr = 0;}
        if (!fileAccessMap.containsKey(modelnum)) {
            modelnr++;
            fileAccessMap.put(modelnum,null);
        }
        saveModel();
    }

    public void addModel(Structure model) {
        currentModel = model;
        int modelnum = ++currentModelnr;

        addModel(model,modelnum);
   }

    public void toPDB(String baseFN) {
        setBaseFileName(baseFN);
        toPDB();
    }

    public void toPDB() {

        System.out.println("Saving structure with "+modifiedModelMap.size()+"/"+modelnr);
        saveModel();
        Structure start = getModel(0);
        if (headerContainer.getHeader() != null) start.setHeader(headerContainer.getHeader());
        start.setNmr(true);
        String fileName = baseFileName+".pdb.gz";
        try {
            FileOutputStream dest = new FileOutputStream(fileName);
            GZIPOutputStream out = new GZIPOutputStream(new PrintStream(dest));

            PrintStream ps = new PrintStream(out);
            ps.println(start.toPDB().trim());

            currentModelnr = 1;
            while(hasNextModel()) {
                System.out.print("Saving Model "+currentModelnr+"/"+modelnr);
                Structure model = nextModel();
                if (model == null) continue;

                Integer cmni = currentModelnr+1;
                String cmn = cmni.toString();
                String cmns = "";
                for (int i = 0; i < 4-cmn.length(); i++) cmns += " ";
                ps.println("MODEL      "+cmns+cmn);
                ps.println(model.toPDB().trim());

                (new File(modifiedModelMap.get(currentModelnr)) ).delete();
            }
            out.close();
        }
        catch (Exception e) {
            System.err.println("Error saving structure");
            e.printStackTrace();
        }

    }

}